diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__init__.py b/lib/python3.12/site-packages/torch/_higher_order_ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d235ee11d064d057c56a5839893ecbb1ab8ac105 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/__init__.py @@ -0,0 +1,22 @@ +from torch._higher_order_ops.cond import cond +from torch._higher_order_ops.flex_attention import ( + flex_attention, + flex_attention_backward, +) +from torch._higher_order_ops.hints_wrap import hints_wrapper +from torch._higher_order_ops.invoke_subgraph import invoke_subgraph +from torch._higher_order_ops.prim_hop_base import PrimHOPBase +from torch._higher_order_ops.scan import scan +from torch._higher_order_ops.while_loop import while_loop + + +__all__ = [ + "cond", + "while_loop", + "invoke_subgraph", + "scan", + "flex_attention", + "flex_attention_backward", + "hints_wrapper", + "PrimHOPBase", +] diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0418c4d993c5fccad39eb7532ca6134b6d202f0 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/associative_scan.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/associative_scan.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c2a693c651d18e1cdb3f81706f17f4b2269598d Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/associative_scan.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/auto_functionalize.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/auto_functionalize.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d91178a8940cb03ff35b247af84e9d8920a9120d Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/auto_functionalize.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/cond.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/cond.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d196844caa749a411ac4c2b3e50662de9fcc44ef Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/cond.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/effects.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/effects.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..415d4b5e1ae51ca7d1c5e9fafe544544db06b75b Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/effects.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/executorch_call_delegate.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/executorch_call_delegate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1dc3c7091da881285491f88fac57f19272846e78 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/executorch_call_delegate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/flex_attention.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/flex_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a19590a768bd6440127d1f8956c70ba299d9842c Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/flex_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/hints_wrap.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/hints_wrap.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df26deb4b7536456fe6e430419a69d0df7ac7dcd Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/hints_wrap.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/invoke_subgraph.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/invoke_subgraph.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86404fb1371a5e90a40faf8495ea695892b9bc73 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/invoke_subgraph.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/map.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/map.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4781a8f7b33cb788fdcbca99e6e432fd9d8b632e Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/map.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/out_dtype.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/out_dtype.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb9c4f59fa45666e2063748f0ace5bb5947b56cc Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/out_dtype.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/prim_hop_base.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/prim_hop_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..254bac9c6be4512ae5a00aa5d7212071d7ee5dbc Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/prim_hop_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/run_const_graph.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/run_const_graph.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..998ee65db38f03a115e32ca92a62f26f82e13dd7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/run_const_graph.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/scan.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/scan.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32630b671074fb1b29867d57e3d3564f83a0d877 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/scan.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/strict_mode.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/strict_mode.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76a558981c9c86c1945a5d9908801709379575e4 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/strict_mode.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/torchbind.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/torchbind.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d5b92a35112ef4d77b8bcb8f8a9ae8c174be423e Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/torchbind.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/triton_kernel_wrap.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/triton_kernel_wrap.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4b2f3ceeda58cd754d1a647008b3ffa40a67772 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/triton_kernel_wrap.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58b109c55b42035074234a6c54eb8794aa33c92d Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/while_loop.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/while_loop.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e5158ee272538f8e25a477cfef183e314e05e27 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/while_loop.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/wrap.cpython-312.pyc b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/wrap.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da1021b65127c95c575f739717186991c7cd4ad8 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_higher_order_ops/__pycache__/wrap.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/associative_scan.py b/lib/python3.12/site-packages/torch/_higher_order_ops/associative_scan.py new file mode 100644 index 0000000000000000000000000000000000000000..c1db33dfa34aaefbc85a62223111fee4f1c131e9 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/associative_scan.py @@ -0,0 +1,441 @@ +# mypy: allow-untyped-defs +import functools +import itertools +from typing import Any, Callable, List + +import torch +import torch._prims_common as utils +import torch._subclasses.functional_tensor +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._higher_order_ops.utils import ( + _maybe_run_with_interpreter, + _set_compilation_env, + autograd_not_implemented, + first_slice_copy, + reenter_make_fx, + unique_graph_id, +) +from torch._inductor.utils import is_pointwise_use +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ProxyTorchDispatchMode, + track_tensor_tree, +) + + +aten = torch._ops.ops.aten + + +def wrap_combine_fn_flat(*args, combine_fn, spec, num_leaves): + assert len(args) == 2 * num_leaves + lhs = pytree.tree_unflatten(args[:num_leaves], spec) + rhs = pytree.tree_unflatten(args[num_leaves:], spec) + combined = combine_fn(lhs, rhs) + combined_leaves = pytree.tree_leaves(combined) + assert num_leaves == len(combined_leaves) + return combined_leaves + + +def _interleave(a, b, dim): + # https://stackoverflow.com/questions/60869537/how-can-i-interleave-5-pytorch-tensors + if b_trunc := (a.shape[dim] == b.shape[dim] + 1): + pad = ( + [0] * ((b.ndim - dim - 1) * 2 + 1) + + [1] + + [0] * (b.ndim * 2 - ((b.ndim - dim - 1) * 2 + 2)) + ) + b = torch.nn.functional.pad(b, pad) + + stacked = torch.stack([a, b], dim=dim + 1) + interleaved = torch.flatten(stacked, start_dim=dim, end_dim=dim + 1) + if b_trunc: + # TODO: find torch alternative for slice_along dim for torch.jit.script to work + interleaved = aten.slice(interleaved, dim, 0, b.shape[dim] + a.shape[dim] - 1) + return interleaved + + +def safe_map(f, *args): + args = list(map(list, args)) + n = len(args[0]) + for arg in args[1:]: + if len(arg) != n: + raise ValueError("length mismatch: {list(map(len, args))}") + + def nf(a): + return f(*a) + + return list(map(nf, zip(*args))) + + +class AssociativeScanOp(HigherOrderOperator): + def __init__(self): + super().__init__("associative_scan") + + def __call__(self, combine_fn, xs, dim): + return super().__call__(combine_fn, xs, dim) + + +associative_scan_op = AssociativeScanOp() + + +def associative_scan( + combine_fn: Callable[[pytree.PyTree, pytree.PyTree], pytree.PyTree], + xs: pytree.PyTree, + dim: int, + reverse: bool = False, + combine_mode: str = "pointwise", +) -> torch.Tensor: + r""" + Performs an inclusive scan with an associative combine function. + + .. warning:: + `torch.associative_scan` is a prototype feature in PyTorch. It currently + does not support autograd and you may run into miscompiles. + Read more about feature classification at: + https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype + + This operator requires runtime code generation and so requires support for + ``torch.compile``. Further, only CUDA device codegen is supported at the moment. + + Args: + combine_fn (Callable): A binary callable with type ``(Tensor, Tensor) -> Tensor``, + or if input is a pytree ``(pytree, pytree) -> pytree``. + This function must be pure, i.e., no lifted arguments are supported at the moment, + satisfy the associative property and have no side-effects. + xs (torch.Tensor): The input tensor, or nested pytree of tensors. + All inputs are expected to have the same shape. + dim (int): the dimension to scan over + reverse (bool): A boolean stating if the scan should be reversed with respect to ``dim``, default ``False``. + combine_mode (str): A string indicating whether the ``combine_fn`` is ``pointwise`` or ``generic``, default ``pointwise``. + If ``combine_mode=pointwise``, ``combine_fn`` must be pure, may only contain pointwise operations + and ``xs`` must be CUDA tensors. + In all other cases ``combine_mode=generic`` should be used. + Note: ``combine_mode=pointwise`` is more efficient than ``combine_mode=generic``. + + + Example:: + + def add(x: torch.Tensor, y: torch.Tensor): + return x + y + + cumsum = associative_scan(add, x, dim) + + """ + if not callable(combine_fn): + raise ValueError("Combine_fn must be a callable, but got {combine_fn}") + if not isinstance(dim, int): + raise ValueError("Dim must be an int, but got " + str(type(dim))) + if combine_mode not in ["pointwise", "generic"]: + raise ValueError( + "Combine_mode must either 'pointwise' or 'generic', but got {combine_mode}" + ) + + if not torch._dynamo.is_compiling(): + with _set_compilation_env(), torch._dynamo.utils.disable_cache_limit(): + return torch.compile(associative_scan, fullgraph=True)( + combine_fn, xs, dim, reverse=reverse, combine_mode=combine_mode + ) + + leaves, spec = pytree.tree_flatten(xs) + + if combine_mode == "pointwise" and not all(l.device.type == "cuda" for l in leaves): + raise ValueError( + "For combine_mode='pointwise', all input tensors need to be on CUDA" + ) + + if len(leaves) == 0: + raise ValueError("Expected at least 1 xs leaf") + if any(not isinstance(x, torch.Tensor) for x in leaves): + raise ValueError("xs leaves must be a Tensor") + if any(x.is_sparse for x in leaves): + raise ValueError("xs leaves must dense Tensors, consider using `to_dense()`") + if any(x.ndim <= dim for x in leaves): + raise ValueError( + "All xs leaves must at least have 'dim' number of dimensions and scan dimension > 0" + ) + if any(x.shape[dim] == 0 for x in leaves): + raise ValueError( + "All xs leaves must at least have 'dim' number of dimensions and scan dimension > 0" + ) + + if reverse: + leaves = [torch.flip(elem, [dim]) for elem in leaves] + + ndim = leaves[0].ndim + dim = utils.canonicalize_dim(ndim, dim) + + # Call the combine_fn with only a slice along the scan dim + # and check whether the output leaves have the same slice dimensions + sliced_leaves = [first_slice_copy(leaf, dim) for leaf in leaves] + + out = combine_fn( + pytree.tree_unflatten(sliced_leaves, spec), + pytree.tree_unflatten(sliced_leaves, spec), + ) + out_leaves = pytree.tree_leaves(out) + if len(leaves) != len(out_leaves): + raise RuntimeError( + "The number of leaves of the pytree of the output of the operator needs to match the length of the pytree of the input" + ) + if any( + x.shape != x_sliced.shape + or x.dtype != x_sliced.dtype + or x.device != x_sliced.device + or x.stride() != x_sliced.stride() + for x, x_sliced in zip(out_leaves, sliced_leaves) + ): + raise RuntimeError( + f"The metadata of the output of the operator needs to match the meta data of the xs pytree" + f"\n xs metadata : {[(x.shape, x.dtype, x.device, x.stride()) for x in sliced_leaves]}" + f"\n operator output metadata: {[(x.shape, x.dtype, x.device, x.stride()) for x in out_leaves]}" + ) + + if combine_mode == "generic": + # The generic_associative_scan implementation calls the combine_fn with a `batch` along the scan dimension + # For example, consider: + # def add(x: torch.Tensor, y: torch.Tensor): + # return x + y + # leaves = torch.tensor([[0.0, 1.0, 2.0, 3.0] + # [0.0, 1.0, 2.0, 3.0]]) + # which has shape 2 x 4; + # dim = 1; + # In the first iteration of `_scan` the combine_fn gets invoked with + # combine_fn([torch.tensor([[0.0, 2.0], + # [0.0, 2.0]])], + # [torch.tensor([[1.0, 3.0], + # [1.0, 3.0]])]) + # The arguments are of shape 2 x 2, but can be evaluated in parallel along the scan dimension. + # TODO: In case of the additional inputs, we the in_dims should be set to None + combine_fn = functools.partial( + wrap_combine_fn_flat, + combine_fn=torch.vmap( + combine_fn, + in_dims=( + pytree.tree_unflatten([dim] * len(leaves), spec), + pytree.tree_unflatten([dim] * len(leaves), spec), + ), + out_dims=dim, + ), + spec=spec, + num_leaves=len(leaves), + ) + result_flat = generic_associative_scan(combine_fn, leaves, dim) + else: + combine_fn = functools.partial( + wrap_combine_fn_flat, + combine_fn=combine_fn, + spec=spec, + num_leaves=len(leaves), + ) + result_flat = associative_scan_op(combine_fn, leaves, dim) + + if reverse: + result_flat = [torch.flip(elem, [dim]) for elem in result_flat] + + return pytree.tree_unflatten(result_flat, spec) + + +def generic_associative_scan(operator, leaves, dim=0): + r""" + This function performs the associative_scan operation. + The algorithm works by recursively collecting neighbours of ``leaves`` and subsequently + applying the ``operator`` on all pairs in parallel along ``dim``. + The results of the recursive calls are later combined. + + Args: + operator (Callable): A binary callable with type ``(Tensor, Tensor) -> Tensor``, + or if input is a pytree ``(pytree, pytree) -> pytree``. + This function must be pure, pointwise, and satisfy the associative property. + leaves (torch.Tensor): A list of torch.Tensors converted from the pytree of + ``xs`` provided to ``associative_scan``. + All inputs are expected to have the same shape. + dim (int): the dimension to scan over + + + Example:: + + def add(x: torch.Tensor, y: torch.Tensor): + return x + y + + leaves = torch.tensor([0.0, 1.0, 2.0, 3.0]) + + First iteration of _scan -> + # odd_elems -> apply operator on all neighbours + # odd_elems = operator([torch.tensor([0.0, 2.0])], + # [torch.tensor([1.0, 3.0])]) + odd_elems = torch.tensor([1.0, 5.0]) + Second iteration of _scan -> + # odd_elems = operator([torch.tensor([1.0])], + # [torch.tensor([5.0])]) + odd_elems = torch.tensor([6.0]) + # even_elems -> apply operator on all odd_elems and + # every second element of ``elems``, starting from the second element. + # even_elems is expanded with the first element of ``elems`` + even_elems = [1.0] + # Merges odd_elems and even_elems + res = torch.tensor([1.0, 6.0]) + # even_elems -> apply operator on all odd_elems and + # every second element of ``elems``, starting from the second element. + # even_elems is expanded with the first element of ``elems`` + even_elems = [0.0, 3.0] + # Merges odd_elems and even_elems + res = torch.tensor([0.0, 1.0, 3.0, 6.0]) + + """ + + def _scan(elems): + """Perform the actual recursive scan on ``elems``.""" + num_elems = elems[0].shape[dim] + + if num_elems < 2: + return elems + + reduced_elems = operator( + *[aten.slice(elem, dim, 0, -1, 2) for elem in elems], + *[aten.slice(elem, dim, 1, None, 2) for elem in elems], + ) + + # Recursively compute scan for partially reduced tensors. + odd_elems = _scan(reduced_elems) + + if num_elems % 2 == 0: + even_elems = operator( + *[aten.slice(e, dim, 0, -1) for e in odd_elems], + *[aten.slice(e, dim, 2, None, 2) for e in elems], + ) + else: + even_elems = operator( + *odd_elems, + *[aten.slice(e, dim, 2, None, 2) for e in elems], + ) + + # The first element of a scan is the same as the first element + # of the original `elems`. + even_elems = [ + torch.cat([aten.slice(elem, dim, 0, 1), result], dim=dim) + if result.shape.numel() > 0 and elem.shape[dim] > 0 + else result + if result.shape.numel() > 0 + else aten.slice( + elem, dim, 0, 1 + ) # Jax allows/ignores concat with 0-dim, Pytorch does not + for (elem, result) in zip(elems, even_elems) + ] + + return list( + safe_map(functools.partial(_interleave, dim=dim), even_elems, odd_elems) + ) + + scans = _scan(leaves) + + return scans + + +def trace_associative_scan( + proxy_mode, func_overload, combine_fn: Callable, xs: List[torch.Tensor], dim: int +): + with disable_proxy_modes_tracing(): + sample_xs = [first_slice_copy(x, dim) for x in itertools.chain(xs, xs)] + combine_graph = reenter_make_fx(combine_fn)(*sample_xs) + + outputs = None + for node in combine_graph.graph.nodes: + if node.op == "output": + assert outputs is None + assert len(node.args) == 1 + outputs = node.args[0] + + if not all(is_pointwise_use(use) or use.op == "output" for use in node.users): + raise ValueError( + "For combine_mode='pointwise', the combine_fn needs to be pointwise" + ) + + assert outputs is not None + assert len(outputs) == len( + xs + ), f"expected combine_fn to return {len(xs)} results but got {len(outputs)}" + + for i, o in zip(xs, outputs): + o_meta = o.meta["tensor_meta"] + assert o_meta.dtype == i.dtype, ( + f"combine_fn output type mismatch, expected {i.dtype} " + + f"but got {o_meta.dtype}" + ) + + _, combine_graph_name = unique_graph_id(proxy_mode, prefix="scan_combine_graph") + + proxy_mode.tracer.root.register_module(combine_graph_name, combine_graph) + + args = (combine_graph, xs, dim) + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, args) + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", func_overload, proxy_args, {}, name="associative_scan" + ) + + with disable_proxy_modes_tracing(): + out = [aten.clone(x) for x in xs] + + return track_tensor_tree(out, out_proxy, constant=None, tracer=proxy_mode.tracer) + + +@associative_scan_op.py_impl(DispatchKey.CompositeExplicitAutograd) +def associative_scan_op_dense(combine_fn, xs, dim): + return generic_associative_scan(combine_fn, xs, dim) + + +associative_scan_op.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(associative_scan_op, deferred_error=True) +) + + +@associative_scan_op.py_impl(ProxyTorchDispatchMode) +def associative_scan_proxy_mode(mode, combine_fn, xs, dim): + return trace_associative_scan(mode, associative_scan_op, combine_fn, xs, dim) + + +@associative_scan_op.py_impl(FakeTensorMode) +def assoiciative_scan_fake_tensor_mode(mode, combine_fn, xs, dim): + with mode: + return [x.clone() for x in xs] + + +@associative_scan_op.py_functionalize_impl +def associative_scan_functionalize(ctx, combine_fn, xs, dim): + unwrapped_xs = ctx.unwrap_tensors(xs) + with ctx.redispatch_to_next() as m: + functional_combine_fn = ctx.functionalize( + _maybe_run_with_interpreter(combine_fn) + ) + ret = associative_scan_op(functional_combine_fn, unwrapped_xs, dim) + return ctx.wrap_tensors(ret) + + +def _fake_associative_scan(combine_fn, xs, dim, reverse=False): # noqa: F811 + inp_leaves, spec = pytree.tree_flatten(xs) + result_flat: List[Any] = [] + num_leaves = len(inp_leaves) + op = reversed if reverse else lambda x: x + + for ind in op(range(inp_leaves[0].size(dim))): + r = [ + inp_leaves[leave_ind][(slice(None),) * dim + (ind,)] + for leave_ind in range(num_leaves) + ] + if (ind > 0 and not reverse) or ( + ind < (inp_leaves[0].size(dim) - 1) and reverse + ): + r = combine_fn( + pytree.tree_unflatten(result_flat[-1], spec), + pytree.tree_unflatten(r, spec), + ) + r_flat, _ = pytree.tree_flatten(r) + result_flat.append(r_flat) + + results = [ + torch.stack([e[leave_ind] for e in op(result_flat)], dim) + for leave_ind in range(num_leaves) + ] + return pytree.tree_unflatten(results, spec) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/auto_functionalize.py b/lib/python3.12/site-packages/torch/_higher_order_ops/auto_functionalize.py new file mode 100644 index 0000000000000000000000000000000000000000..7557deede66d5114551f3fbb6ecba9bbffec24b5 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/auto_functionalize.py @@ -0,0 +1,832 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +import warnings +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union + +import torch +import torch.utils._pytree as pytree +from torch import Tensor +from torch._C import DispatchKey +from torch._ops import HigherOrderOperator, OperatorBase, OpOverload +from torch._prims_common import clone_preserve_strides +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ProxyTorchDispatchMode, + track_tensor_tree, +) + + +def get_base(tensor): + if torch.is_inference_mode_enabled(): + return tensor._inference_mode_base + else: + return tensor._base + + +class ViewInfo(ABC): + base_index: int + + def __init__(self, base_index): + self.base_index = base_index + + @abstractmethod + def regenerate_view(self, bases_list: List[Tensor]): + pass + + +@dataclass +class AsStridedViewInfo(ViewInfo): + size: Sequence[Union[int, torch.SymInt]] + stride: Sequence[Union[int, torch.SymInt]] + storage_offset: int + + def __init__(self, base_index, size, stride, storage_offset): + super().__init__(base_index) + self.size = size + self.stride = stride + self.storage_offset = storage_offset + + def regenerate_view(self, bases_list: List[Tensor]): + return torch.as_strided( + bases_list[self.base_index], + self.size, + self.stride, + self.storage_offset, + ) + + +@dataclass +class SliceViewInfo(ViewInfo): + dim: Union[int, torch.SymInt] + start: Union[int, torch.SymInt] + end: Union[int, torch.SymInt] + + def __init__(self, base_index, dim, start, end): + super().__init__(base_index) + self.dim = dim + self.start = start + self.end = end + + def regenerate_view(self, bases_list: List[Tensor]): + return torch.ops.aten.slice.Tensor( + bases_list[self.base_index], self.dim, self.start, self.end + ) + + +@dataclass +class AliasViewInfo(ViewInfo): + def __init__(self, base_index): + super().__init__(base_index) + + def regenerate_view(self, bases_list: List[Tensor]): + return torch.ops.aten.alias.default(bases_list[self.base_index]) + + +@dataclass +class NotView(ViewInfo): + def __init__(self, base_index): + super().__init__(base_index) + + def regenerate_view(self, bases_list: List[Tensor]): + return bases_list[self.base_index] + + +def is_alias(base, tensor): + from torch.fx.experimental.symbolic_shapes import statically_known_true, sym_eq + + return all( + statically_known_true(a) + for a in [ + sym_eq(base.storage_offset(), tensor.storage_offset()), + sym_eq(base.stride(), tensor.stride()), + sym_eq(base.size(), tensor.size()), + ] + ) + + +# return None or (dim, start, end) +def try_use_slice(base, tensor): + from torch.fx.experimental.symbolic_shapes import statically_known_true, sym_eq + + # This condition should never be triggered. + if is_alias(base, tensor): + return (0, 0, base.size()[0]) + + # TODO is there cases can we use slice even if stride or len(sizes) are not equal? + if not statically_known_true(sym_eq(tensor.stride(), base.stride())): + return None + if not statically_known_true(sym_eq(len(tensor.size()), len(base.size()))): + return None + + dim = None + count = 0 + for i in range(len(tensor.size())): + if base.size()[i] != tensor.size()[i]: + dim = i + count = count + 1 + if count != 1: + return None + + if tensor.storage_offset() % tensor.stride()[dim] != 0: + return None + start = tensor.storage_offset() // tensor.stride()[dim] + end = start + tensor.size()[dim] + return (dim, start, end) + + +def write_view_information_to_args( + mutable_arg_names: List[str], + mutable_arg_types: List[torch.Type], + kwargs: Dict[str, Any], + arg_to_base_index: Dict[str, Any], +): + """ + This function writes the view information into kwargs. It reads mutable_args from kwargs. + and uses arg_to_base_index and tensor information to write ViewInfo into kwargs. + mutable_arg_names: mutable custom operator arg names. + mutable_arg_types: mutable custom operator arg types. + kwargs: the original custom operator args. + arg_to_base_index: maps mutable_arg_name to int | [int] that refers to the base tensor that + corresponds to the input tensor + """ + + def write_single_view(prefix: str, tensor: Tensor, base_index: int): + assert f"{prefix}_base_index" not in kwargs + assert f"{prefix}_size" not in kwargs + assert f"{prefix}_stride" not in kwargs + assert f"{prefix}_storage_offset" not in kwargs + + assert f"{prefix}_slice_dim" not in kwargs + assert f"{prefix}_slice_start" not in kwargs + assert f"{prefix}_slice_end" not in kwargs + + def use_as_strided(tensor): + kwargs[f"{prefix}_size"] = tensor.size() + kwargs[f"{prefix}_stride"] = tensor.stride() + kwargs[f"{prefix}_storage_offset"] = tensor.storage_offset() + + def use_slice(dim, start, end): + kwargs[f"{prefix}_slice_dim"] = dim + kwargs[f"{prefix}_slice_start"] = start + kwargs[f"{prefix}_slice_end"] = end + + def use_alias(): + kwargs[f"{prefix}_alias"] = True + + # The start if the function + if tensor is None: + kwargs[f"{prefix}_base_index"] = None + else: + base = get_base(tensor) + kwargs[f"{prefix}_base_index"] = base_index + if base is None: + # no need to add anything else other than _base_index + return + elif is_alias(base, tensor): + use_alias() + elif (slice_info := try_use_slice(base, tensor)) is not None: + use_slice(*slice_info) + else: + use_as_strided(tensor) + + for arg_name, arg_type in zip(mutable_arg_names, mutable_arg_types): + arg = kwargs[arg_name] + if isinstance(arg_type, torch.ListType): + if arg is None: + kwargs[f"_{arg_name}_length"] = None + + kwargs[f"_{arg_name}_length"] = len(arg) + for i, elem in enumerate(arg): + write_single_view( + f"_{arg_name}_{i}", elem, arg_to_base_index[arg_name][i] + ) + + elif isinstance(arg_type, (torch.TensorType, torch.OptionalType)): + write_single_view( + f"_{arg_name}", + kwargs[arg_name], + arg_to_base_index.get(arg_name, None), + ) + else: + raise RuntimeError(f"Unsupported type {arg_type}") + + +# Returns a dict of arg_name -> ViewInfo | [ViewInfo] +def read_view_information_from_args( + mutable_arg_names: List[str], + mutable_arg_types: List[torch.Type], + kwargs: Dict[str, Any], + all_bases: List[Tensor], +): + """ + This reads the view information added by `write_view_information_to_args` from kwargs, pop them, + and returns a dict arg_name -> ViewInfo | [ViewInfo](if the input is list). that maps each mutable arg + to its view information. + mutable_arg_names: mutable custom operator arg names. + mutable_arg_types: mutable custom operator arg types. + kwargs : args of auto_functionalize(custom_op, kwargs) + """ + + def get_arg(name): + return kwargs.pop(name) + + def read_single_view(prefix): + base_index = get_arg(f"{prefix}_base_index") + if base_index is None: + return None + elif f"{prefix}_alias" in kwargs: + get_arg(f"{prefix}_alias") + return AliasViewInfo(base_index) + elif f"{prefix}_storage_offset" in kwargs: + # The view is regenerated using as_strided. + size = get_arg(f"{prefix}_size") + stride = get_arg(f"{prefix}_stride") + storage_offset = get_arg(f"{prefix}_storage_offset") + return AsStridedViewInfo(base_index, size, stride, storage_offset) + elif f"{prefix}_slice_dim" in kwargs: + dim = get_arg(f"{prefix}_slice_dim") + start = get_arg(f"{prefix}_slice_start") + end = get_arg(f"{prefix}_slice_end") + return SliceViewInfo(base_index, dim, start, end) + else: + # This means that the argument is the base tensor + return NotView(base_index) + + args_view_info: Dict[str, Any] = {} + for arg_name, arg_type in zip(mutable_arg_names, mutable_arg_types): + if isinstance(arg_type, torch.ListType): + length = get_arg(f"_{arg_name}_length") + if length is None: + # The whole list is None. + args_view_info[arg_name] = None + else: + args_view_info[arg_name] = [ + read_single_view(f"_{arg_name}_{i}") for i in range(length) + ] + + elif isinstance(arg_type, (torch.TensorType, torch.OptionalType)): + args_view_info[arg_name] = read_single_view(f"_{arg_name}") + else: + raise RuntimeError(f"Unsupported type {arg_type}") + return args_view_info + + +# NOTE: [auto-functionalizing custom ops] +# Users may wish to torch.compile custom ops that mutate their inputs. +# torch.compile will automatically support this op without anyone needing +# to provide a functionalization kernel for it. Here's how. +# +# Let's say we have a hypothetical mylib::sin_(Tensor(a!) x) -> () +# op. First, when FakeTensor sees this op: +# - If the schema says it returns nothing, we can generate a trivial +# FakeTensor rule for it (that returns nothing). +# - Otherwise, the user needs to provide a FakeTensor impl (fake impl) +# +# Next, when Python FunctionalTensor sees the op, it will functionalize +# it by emitting a call to an auto_functionalize(op, ["x"], {"x": ...}) +# HOP and replacing the mutated inputs with corresponding outputs of this HOP. +# This HOP effectively runs the functional version of the op when +# called: it clones inputs that will be mutated, runs the op, and +# then returns (output, Tensors with the new values) +# +# auto_functionalize_v2 is an improved version of auto_functionalize that better handle +# re-inplacing views. + + +class AutoFunctionalized(HigherOrderOperator): + """auto_functionalized(_mutable_op, **kwargs) + + This HOP runs a "functional" version of _mutable_op. + + Concretely, it looks at all the arguments that are mutable through + _mutable_op's operator schema, clones those kwargs, runs + `out = _mutable_op(**kwargs)` with the cloned values, and then returns the + operator output concatenated with the cloned values that were mutated. + + We have some restrictions on `_mutable_op`. + See `can_auto_functionalize` for the restrictions. We can likely lift + many of these if users request it. + + The reason why _mutable_op is prefixed with an + underscore is to prevent collisions with kwarg names in **kwargs. + """ + + def __init__(self) -> None: + super().__init__("auto_functionalized") + + def __call__( + self, + /, + _mutable_op: OpOverload, + **kwargs: Any, + ) -> Tuple[Any, Tuple[Tensor, ...]]: + assert can_auto_functionalize(_mutable_op) + assert isinstance(kwargs, dict) + return super().__call__(_mutable_op, **kwargs) + + +auto_functionalized = AutoFunctionalized() +auto_functionalized.__module__ = "torch.ops.higher_order" + +auto_functionalized.fallthrough(DispatchKey.AutogradCPU) +auto_functionalized.fallthrough(DispatchKey.AutogradCUDA) + + +class AutoFunctionalizedV2(HigherOrderOperator): + """auto_functionalized_v2(_mutable_op, **kwargs) + + This HOP runs a "functional" version of _mutable_op. + Unlike AutoFunctionalized, this version is improved to better handle + view tensors. This version is only used in non export mode. + """ + + def __init__(self) -> None: + super().__init__("auto_functionalized_v2") + + def __call__( + self, + /, + _mutable_op: OpOverload, + **kwargs: Any, + ) -> Tuple[Any, Tuple[Tensor, ...]]: + assert can_auto_functionalize(_mutable_op) + assert isinstance(kwargs, dict) + return super().__call__(_mutable_op, **kwargs) + + +auto_functionalized_v2 = AutoFunctionalizedV2() +auto_functionalized_v2.__module__ = "torch.ops.higher_order" + +auto_functionalized_v2.fallthrough(DispatchKey.AutogradCPU) +auto_functionalized_v2.fallthrough(DispatchKey.AutogradCUDA) + + +def can_auto_functionalize(op: OperatorBase) -> bool: + if not isinstance(op, OpOverload): + return False + + if torch._library.utils.is_builtin(op): + # We control the built-ins. These may (in rare cases) + # do input metadata mutation (which we have banned on custom ops) + return False + schema = op._schema + if not schema.is_mutable: + return False + schema = op._schema + + for arg in schema.arguments: + if arg.alias_info is None: + continue + if not arg.alias_info.is_write: + continue + if type(arg.type) is torch.TensorType: + continue + if ( + type(arg.type) is torch.OptionalType + and type(arg.type.getElementType()) is torch.TensorType + ): + continue + if ( + type(arg.type) is torch.ListType + and type(arg.type.getElementType()) is torch.TensorType + ): + continue + # Not yet supported: other Tensor types. This includes things like + # Tensor?[], Tensor[]?. + return False + + if len(schema.returns) == 1 and isinstance(schema.returns[0].type, torch.NoneType): + # Skip schema returns -> None + return True + # The returns must not alias anything + for ret in schema.returns: + if ret.alias_info is None and type(ret.type) is torch.TensorType: + continue + # Not yet supported: List[Tensor] return. + return False + if torch._C._dispatch_has_kernel_for_dispatch_key(op.name(), "Functionalize"): + return False + return True + + +def get_mutable_args(op: OpOverload) -> Tuple[List[str], List[torch.Type]]: + """ + Returns the list of argument names that get mutated according to the + schema and their types. + """ + mutable_args_names = [ + arg.name + for arg in op._schema.arguments + if arg.alias_info is not None and arg.alias_info.is_write + ] + + mutable_args_types = [ + arg.type + for arg in op._schema.arguments + if arg.alias_info is not None and arg.alias_info.is_write + ] + return mutable_args_names, mutable_args_types + + +def do_auto_functionalize( + op: OpOverload, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], +) -> Any: + """Functionalizes a call to op(*args, **kwargs) by emitting a call to + `outs = auto_functionalized(op, normalized_kwargs)` + and replacing the mutated (args, kwargs) with the corresponding outputs. + + The normalized_kwargs are just the (args, kwargs), but all in kwarg form. + This makes handling easier for the auto_functionalized HOP. + """ + from torch._subclasses.functional_tensor import PythonFunctionalizeAPI + + ctx = PythonFunctionalizeAPI() + + # All of the (args, kwargs), but all as kwargs. The names for the + # args come from the schema. This makes it easier for us to work with them. + normalized_kwargs = {} + schema = op._schema + for idx, arg in enumerate(schema.arguments): + # NB: torch_dispatch kwargs are the args defined as kwarg-only in the schema + if arg.name in kwargs: + normalized_kwargs[arg.name] = kwargs[arg.name] + elif idx < len(args): + # if its out of bounds we don't need to do anything + # as it means the the optional arg was passed with its default + # value + normalized_kwargs[arg.name] = args[idx] + else: + normalized_kwargs[arg.name] = arg.default_value + + unwrapped_kwargs = ctx.unwrap_tensors(normalized_kwargs) # type: ignore[arg-type] + if "self" in unwrapped_kwargs or "self_" in unwrapped_kwargs: + warnings.warn( + "Using `self` or `self_` as an argument in the definition of custom ops may lead to ambiguous parsing. " + "Please consider using a different name for this argument to avoid potential issues." + ) + with ctx.redispatch_to_next(): + unwrapped_outs = auto_functionalized( + op, **unwrapped_kwargs # type: ignore[arg-type] + ) + + # List of the name of args that get mutated (according to the schema) + mutable_args_names, _ = get_mutable_args(op) + + unwrapped_actual_out: Union[Any, Tuple[Any]] = unwrapped_outs[ + : -len(mutable_args_names) + ] + unwrapped_mutable_out = unwrapped_outs[-len(mutable_args_names) :] + + if len(op._schema.returns) == 0: + assert unwrapped_actual_out[0] is None + unwrapped_actual_out = None + elif len(op._schema.returns) == 1: + assert len(unwrapped_actual_out) == 1 + unwrapped_actual_out = unwrapped_actual_out[0] + else: + assert len(unwrapped_actual_out) == len(op._schema.returns) + + for name, unwrapped_out in zip(mutable_args_names, unwrapped_mutable_out): + # Can be None if input was `Tensor(a!)?` + if unwrapped_out is None: + continue + + # We only handle Tensor or List[Tensor] here for now. + def sync_update(o, orig_arg): + ctx.replace(orig_arg, o) + ctx.commit_update(orig_arg) + ctx.sync(orig_arg) + + orig_arg = normalized_kwargs[name] + + if isinstance(unwrapped_out, torch.Tensor): + sync_update(unwrapped_out, orig_arg) + elif isinstance(unwrapped_out, list) and all( + isinstance(o, torch.Tensor) for o in unwrapped_out + ): + assert len(orig_arg) == len(unwrapped_out) + for orig_a, o in zip(orig_arg, unwrapped_out): + sync_update(o, orig_a) + else: + raise RuntimeError( + f"unsupported type for auto-functionalization: {unwrapped_out}" + ) + + return ctx.wrap_tensors(unwrapped_actual_out) # type: ignore[arg-type] + + +def do_auto_functionalize_v2( + op: OpOverload, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], +) -> Any: + from torch._subclasses.functional_tensor import PythonFunctionalizeAPI + + ctx = PythonFunctionalizeAPI() + + # All of the (args, kwargs), but all as kwargs. The names for the + # args come from the schema. This makes it easier for us to work with them. + normalized_kwargs = {} + + schema = op._schema + for idx, arg in enumerate(schema.arguments): + # NB: torch_dispatch kwargs are the args defined as kwarg-only in the schema + if arg.name in kwargs: + normalized_kwargs[arg.name] = kwargs[arg.name] + elif idx < len(args): + # if its out of bounds we don't need to do anything + # as it means the the optional arg was passed with its default + # value + normalized_kwargs[arg.name] = args[idx] + else: + normalized_kwargs[arg.name] = arg.default_value + + # List of the name of args that get mutated (according to the schema) + mutable_args_names, mutable_args_types = get_mutable_args(op) + + # A list of all bases of mutable args without duplication + all_bases = [] + all_bases_addresses: list[int] = [] + + # Map arg_name to the index of its base in all_bases. + arg_to_base_index: Dict[str, Any] = {} + + def update_dict(tensor, arg_name, index=None): + base = tensor if get_base(tensor) is None else get_base(tensor) + + def set_result(base_index): + if index is None: + arg_to_base_index[arg_name] = base_index + else: + arg_to_base_index[arg_name][index] = base_index + + if not all_bases_addresses.__contains__(base._cdata): + all_bases_addresses.append(base._cdata) + all_bases.append(base) + set_result(len(all_bases) - 1) + else: + set_result(all_bases_addresses.index(base._cdata)) + + for arg_name in mutable_args_names: + arg = normalized_kwargs[arg_name] + if arg is None: + continue + + if isinstance(arg, list): + arg_to_base_index[arg_name] = {} + for i, tensor in enumerate(arg): + if tensor is None: + arg_to_base_index[arg_name].append(None) + continue + + update_dict(tensor, arg_name, i) + + else: + update_dict(arg, arg_name) + + # add view_meta for each args into unwrapped_kwargs. + write_view_information_to_args( + mutable_args_names, + mutable_args_types, + normalized_kwargs, + arg_to_base_index, + ) + + # remove mutated args from the kwargs (its a function of _all_bases now) + for arg_name in mutable_args_names: + del normalized_kwargs[arg_name] # type: ignore[arg-type] + + unwrapped_kwargs = ctx.unwrap_tensors(normalized_kwargs) # type: ignore[arg-type] + if "self" in unwrapped_kwargs or "self_" in unwrapped_kwargs: + warnings.warn( + "Using `self` or `self_` as an argument in the definition of custom ops may lead to ambiguous parsing. " + "Please consider using a different name for this argument to avoid potential issues." + ) + all_basis_unwrapped = ctx.unwrap_tensors(all_bases) + + with ctx.redispatch_to_next(): + unwrapped_outs = auto_functionalized_v2( + op, **dict(unwrapped_kwargs, _all_bases=all_basis_unwrapped) # type: ignore[arg-type] + ) + + unwrapped_actual_out: Union[Any, Tuple[Any]] = ( + unwrapped_outs if len(all_bases) == 0 else unwrapped_outs[: -len(all_bases)] + ) + + unwrapped_mutable_out = ( + [] if len(all_bases) == 0 else unwrapped_outs[-len(all_bases) :] + ) + + if len(op._schema.returns) == 0: + assert unwrapped_actual_out[0] is None + unwrapped_actual_out = None + elif len(op._schema.returns) == 1: + assert len(unwrapped_actual_out) == 1 + unwrapped_actual_out = unwrapped_actual_out[0] + else: + assert len(unwrapped_actual_out) == len(op._schema.returns) + + for orig_arg, unwrapped_out in zip(all_bases, unwrapped_mutable_out): + # Can be None if input was `Tensor(a!)?` + if unwrapped_out is None: + continue + + # We only handle Tensor or List[Tensor] here for now. + def sync_update(o, orig_arg): + ctx.replace(orig_arg, o) + ctx.commit_update(orig_arg) + ctx.sync(orig_arg) + + if isinstance(unwrapped_out, torch.Tensor): + sync_update(unwrapped_out, orig_arg) + elif isinstance(unwrapped_out, list) and all( + isinstance(o, torch.Tensor) for o in unwrapped_out + ): + assert len(orig_arg) == len(unwrapped_out) + for orig_a, o in zip(orig_arg, unwrapped_out): + sync_update(o, orig_a) + else: + raise RuntimeError( + f"unsupported type for auto-functionalization: {unwrapped_out}" + ) + + return ctx.wrap_tensors(unwrapped_actual_out) # type: ignore[arg-type] + + +# auto_functionalize functions +@auto_functionalized.py_impl(DispatchKey.CompositeExplicitAutograd) +def auto_functionalized_dense( + _mutable_op: OpOverload, + _only_clone_these_tensors: Optional[Tuple[str, ...]] = None, + **kwargs: Any, +) -> Tuple[Any, Tuple[Tensor, ...]]: + new_kwargs = dict(**kwargs) + result = [] + + _mutable_args_names, _ = get_mutable_args(_mutable_op) + for name in _mutable_args_names: + if ( + _only_clone_these_tensors is not None + and name not in _only_clone_these_tensors + ): + new_kwargs[name] = kwargs[name] + else: + new_kwargs[name] = ( + [clone_preserve_strides(x) for x in kwargs[name]] + if kwargs[name] is not None and isinstance(kwargs[name], list) + else ( + clone_preserve_strides(kwargs[name]) + if kwargs[name] is not None + else None + ) + ) + result.append(new_kwargs[name]) + out = _mutable_op(**new_kwargs) + + if isinstance(out, tuple): + return (*out, *result) # type: ignore[return-value] + else: + return (out, *result) # type: ignore[return-value] + + +@auto_functionalized.py_impl(FakeTensorMode) +def auto_functionalized_fake( + mode, + _mutable_op: OpOverload, + **kwargs: Any, +) -> Tuple[Any, Tuple[Tensor, ...]]: + with mode: + result = auto_functionalized_dense( + _mutable_op, _only_clone_these_tensors=None, **kwargs + ) + return result + + +@auto_functionalized.py_impl(ProxyTorchDispatchMode) +def auto_functionalized_proxy( + mode, + _mutable_op: OpOverload, + **kwargs: Any, +) -> Tuple[Any, Tuple[Tensor, ...]]: + with disable_proxy_modes_tracing(): + out = auto_functionalized(_mutable_op, **kwargs) + + proxy_kwargs = pytree.tree_map(mode.tracer.unwrap_proxy, kwargs) + out_proxy = mode.tracer.create_proxy( + "call_function", + auto_functionalized, + (_mutable_op,), + proxy_kwargs, + ) + result = track_tensor_tree(out, out_proxy, constant=None, tracer=mode.tracer) + return result + + +@auto_functionalized.py_functionalize_impl +def auto_functionalized_func(ctx, _mutable_op, **kwargs): + unwrapped_kwargs = ctx.unwrap_tensors(kwargs) + with ctx.redispatch_to_next(): + result = auto_functionalized(_mutable_op, **unwrapped_kwargs) + return ctx.wrap_tensors(result) + + +# auto_functionalized_v2 functions +@auto_functionalized_v2.py_impl(DispatchKey.CompositeExplicitAutograd) +def auto_functionalized_v2_dense( + _mutable_op: OpOverload, + _only_clone_these_bases: Optional[Tuple[int, ...]] = None, + **kwargs: Any, +) -> Tuple[Any, Tuple[Tensor, ...]]: + all_bases: List[Tensor] = kwargs.pop("_all_bases", []) + mutable_args_names, mutable_args_types = get_mutable_args(_mutable_op) + args_view_info = read_view_information_from_args( + mutable_args_names, mutable_args_types, kwargs, all_bases + ) + + if _only_clone_these_bases is None: + _only_clone_these_bases = tuple(range(len(all_bases))) + + def maybe_copy(i, t): + if t is None: + return None + if i in _only_clone_these_bases: + return clone_preserve_strides(t) + else: + return t + + all_bases_new = [maybe_copy(i, t) for i, t in enumerate(all_bases)] + + # create new args + new_kwargs = dict(**kwargs) + + # re-generate all inputs from all_bases_new using args_view_info and add them to new_kwargs. + for arg_name in mutable_args_names: + if args_view_info[arg_name] is None: + new_kwargs[arg_name] = None + elif isinstance(args_view_info[arg_name], list): + new_kwargs[arg_name] = [] + for i, elem in enumerate(args_view_info[arg_name]): + if elem is None: + new_kwargs[arg_name].append(None) + else: + view_info = args_view_info[arg_name][i] + new_kwargs[arg_name].append( + view_info.regenerate_view(all_bases_new) + ) + else: + new_kwargs[arg_name] = args_view_info[arg_name].regenerate_view( + all_bases_new + ) + + out = _mutable_op(**new_kwargs) + + if isinstance(out, tuple): + return (*out, *all_bases_new) # type: ignore[return-value] + else: + return (out, *all_bases_new) # type: ignore[return-value] + + +@auto_functionalized_v2.py_impl(FakeTensorMode) +def auto_functionalized_v2_fake( + mode, + _mutable_op: OpOverload, + **kwargs: Dict[str, Any], +) -> Tuple[Any, Tuple[Tensor, ...]]: + with mode: + result = auto_functionalized_v2_dense( + _mutable_op, _only_clone_these_bases=None, **kwargs + ) + return result + + +@auto_functionalized_v2.py_impl(ProxyTorchDispatchMode) +def auto_functionalized_v2_proxy( + mode, + _mutable_op: OpOverload, + **kwargs: Dict[str, Any], +) -> Tuple[Any, Tuple[Tensor, ...]]: + with disable_proxy_modes_tracing(): + out = auto_functionalized_v2(_mutable_op, **kwargs) + + proxy_kwargs = pytree.tree_map(mode.tracer.unwrap_proxy, kwargs) + out_proxy = mode.tracer.create_proxy( + "call_function", + auto_functionalized_v2, + (_mutable_op,), + proxy_kwargs, + ) + result = track_tensor_tree(out, out_proxy, constant=None, tracer=mode.tracer) + return result + + +@auto_functionalized_v2.py_functionalize_impl +def auto_functionalized_v2_func(ctx, _mutable_op, **kwargs): + unwrapped_kwargs = ctx.unwrap_tensors(kwargs) + with ctx.redispatch_to_next(): + result = auto_functionalized_v2(_mutable_op, **unwrapped_kwargs) + return ctx.wrap_tensors(result) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/cond.py b/lib/python3.12/site-packages/torch/_higher_order_ops/cond.py new file mode 100644 index 0000000000000000000000000000000000000000..d5b8c2b5cd2b84c50d53b368d9a065aa817f89db --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/cond.py @@ -0,0 +1,559 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +import contextlib +import logging +import warnings +from typing import Any, Callable, List, Tuple, Union + +import torch +import torch._subclasses.functional_tensor +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._C._functorch import ( + _add_batch_dim, + get_unwrapped, + is_batchedtensor, + maybe_get_bdim, +) +from torch._dispatch.python import suspend_functionalization +from torch._functorch.utils import exposed_in +from torch._guards import detect_fake_mode +from torch._higher_order_ops.utils import ( + _has_potential_branch_input_alias, + _has_potential_branch_input_mutation, + _maybe_run_with_interpreter, + _set_compilation_env, + reenter_make_fx, + save_tensors_and_symints_for_backward, + saved_tensors_and_symints, + unique_graph_id, + UnsupportedAliasMutationException, + validate_subgraph_args_types, +) +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch._subclasses.functional_tensor import disable_functional_mode +from torch.fx.experimental.proxy_tensor import ( + _temp_remove_metadata_torch_function_mode, + _temp_remove_pre_dispatch_torch_function_mode, + disable_proxy_modes_tracing, + ProxyTorchDispatchMode, + track_tensor_tree, +) +from torch.fx.passes.shape_prop import _extract_tensor_metadata +from torch.utils._python_dispatch import _get_current_dispatch_mode + +from .utils import _from_fun, create_fw_bw_graph + + +log = logging.getLogger(__name__) + +""" +We're going to define a `cond_op` operation. +In order to do this, we need implementations for each of the dispatch keys. +""" + + +class CondOp(HigherOrderOperator): + def __init__(self): + super().__init__("cond") + + def __call__(self, pred, true_fn, false_fn, operands): + validate_subgraph_args_types(operands) + return super().__call__(pred, true_fn, false_fn, operands) + + +cond_op = CondOp() + + +@exposed_in("torch") +def cond( + pred: Union[bool, int, float, torch.Tensor], + true_fn: Callable, + false_fn: Callable, + operands: Union[Tuple, List] = (), +) -> Any: + r""" + Conditionally applies `true_fn` or `false_fn`. + + .. warning:: + `torch.cond` is a prototype feature in PyTorch. It has limited support for input and output types and + doesn't support training currently. Please look forward to a more stable implementation in a future version of PyTorch. + Read more about feature classification at: https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype + + `cond` is structured control flow operator. That is, it is like a Python if-statement, + but has restrictions on `true_fn`, `false_fn`, and `operands` that enable it to be + capturable using torch.compile and torch.export. + + Assuming the constraints on `cond`'s arguments are met, `cond` is equivalent to the following:: + + def cond(pred, true_branch, false_branch, operands): + if pred: + return true_branch(*operands) + else: + return false_branch(*operands) + + Args: + pred (Union[bool, torch.Tensor]): A boolean expression or a tensor with one element, + indicating which branch function to apply. + + true_fn (Callable): A callable function (a -> b) that is within the + scope that is being traced. + + false_fn (Callable): A callable function (a -> b) that is within the + scope that is being traced. The true branch and false branch must + have consistent input and outputs, meaning the inputs have to be + the same, and the outputs have to be the same type and shape. + + operands (Tuple of possibly nested dict/list/tuple of torch.Tensor): A tuple of inputs to the + true/false functions. It can be empty if true_fn/false_fn doesn't require input. Defaults to (). + + Example:: + + def true_fn(x: torch.Tensor): + return x.cos() + def false_fn(x: torch.Tensor): + return x.sin() + return cond(x.shape[0] > 4, true_fn, false_fn, (x,)) + + Restrictions: + - The conditional statement (aka `pred`) must meet one of the following constraints: + + - It's a `torch.Tensor` with only one element, and torch.bool dtype + + - It's a boolean expression, e.g. `x.shape[0] > 10` or `x.dim() > 1 and x.shape[1] > 10` + + - The branch function (aka `true_fn`/`false_fn`) must meet all of the following constraints: + + - The function signature must match with operands. + + - The function must return a tensor with the same metadata, e.g. shape, + dtype, etc. + + - The function cannot have in-place mutations on inputs or global variables. + (Note: in-place tensor operations such as `add_` for intermediate results + are allowed in a branch) + + .. warning:: + Temporal Limitations: + + - The **output** of branches must be a **single Tensor**. Pytree of tensors will be supported in the future. + + """ + if torch.compiler.is_dynamo_compiling(): + return cond_op(pred, true_fn, false_fn, operands) + + from torch._dynamo.backends.debugging import ( + make_eager_backend_with_torch_function_mode, + ) + + if isinstance(pred, (bool, int, float)): + # This is the non-strict export case. Strict export and torch.compile are + # handled above in dynamo. + if torch.compiler.is_compiling(): + warnings.warn( + "Pred is a Python constant. When used with torch.cond, it specializes on one of the branches." + " If you want torch.cond to preserve two branches, please make the predicate a boolean tensor or a SymBool.", + UserWarning, + ) + # This is the eager case. We can just run the true or false branch. + if pred: + return true_fn(*operands) + else: + return false_fn(*operands) + + def _validate_input(pred, true_fn, false_fn, operands): + if not isinstance(pred, (bool, torch.Tensor, torch.SymBool)): + raise RuntimeError(f"Expected pred to be bool or tensor, but got {pred}.") + + if isinstance(pred, torch.Tensor) and pred.numel() != 1: + raise RuntimeError( + f"Expected pred to be bool or single-element tensor, but got {pred}." + ) + + if not callable(true_fn) or not callable(false_fn): + raise RuntimeError("Expect both branches to be callable.") + + if not isinstance(operands, (tuple, list)) or pytree.tree_any( + lambda t: not isinstance(t, torch.Tensor), operands + ): + raise RuntimeError( + "Expect operands to be a tuple of possibly nested dict/list/tuple that only " + f"consists of tensor leaves, but got {operands}." + ) + + _validate_input(pred, true_fn, false_fn, operands) + + if not torch._dynamo.is_dynamo_supported(): + raise RuntimeError("torch.cond requires dynamo support.") + + # Dynamo is expecting a callable with "__code__" attribute. + # We cannot directly pass cond_op to it. So we wrap it in a dummy function. + def _cond_op_wrapper(*args, **kwargs): + return cond_op(*args, **kwargs) + + with _set_compilation_env(), torch._dynamo.utils.disable_cache_limit(), _temp_remove_pre_dispatch_torch_function_mode(): + with _temp_remove_metadata_torch_function_mode() as metadata_mode: + if metadata_mode: + backend = make_eager_backend_with_torch_function_mode(metadata_mode) + else: + backend = "eager" + return torch.compile(_cond_op_wrapper, backend=backend, fullgraph=True)( + pred, true_fn, false_fn, operands + ) + + +def create_fw_bw_graph_branches(true_fn, false_fn, *operands): + # See Note [HOP create fw_bw graph] in create_fw_bw_graph in utils.py + + with suspend_functionalization(), disable_functional_mode(): + with disable_proxy_modes_tracing(): + fw_inputs = pytree.tree_map(_from_fun, operands) + + fw_outputs_true = pytree.tree_map(_from_fun, true_fn(*fw_inputs)) + if any( + not isinstance(out, torch.Tensor) + for out in fw_outputs_true + if out is not None + ): + raise RuntimeError( + "Expect outputs of true_fn to only contains tensors or None. " + f"Got types {[type(out) for out in fw_outputs_true]}." + ) + fw_outputs_false = pytree.tree_map(_from_fun, false_fn(*fw_inputs)) + if any( + not isinstance(out, torch.Tensor) + for out in fw_outputs_false + if out is not None + ): + raise RuntimeError( + "Expect outputs of false_fn to only contains tensors or None. " + f"Got types {[type(out) for out in fw_outputs_false]}." + ) + + # TODO: There is a major issue that the create_fw_bw in the higher_order_op is invoked twice: + # Once in the forward path (as it should) and once in the backward path, where it shouldn't be called + # If we can get rid of the second invokation, it would simplify this function + fw_true_graph, joint_true_graph = create_fw_bw_graph( + true_fn, False, fw_inputs, fw_outputs_true + ) + fw_false_graph, joint_false_graph = create_fw_bw_graph( + false_fn, False, fw_inputs, fw_outputs_false + ) + + return fw_true_graph, fw_false_graph, joint_true_graph, joint_false_graph + + +def trace_cond(proxy_mode, func_overload, pred, true_fn, false_fn, operands): + assert isinstance( + operands, (list, tuple) + ), f"Cond operands must be a list or tuple of tensors and SymInts {operands}" + + true_graph = reenter_make_fx(true_fn)(*operands) + false_graph = reenter_make_fx(false_fn)(*operands) + + true_outs = [] + false_outs = [] + for node in true_graph.graph.nodes: + if node.op == "output": + true_outs.extend(node.args) + + for node in false_graph.graph.nodes: + if node.op == "output": + false_outs.extend(node.args) + + flat_true_outs = pytree.arg_tree_leaves(*true_outs) + flat_false_outs = pytree.arg_tree_leaves(*false_outs) + if len(flat_true_outs) != len(flat_false_outs): + raise torch._dynamo.exc.CondOpArgsMismatchError( + f"Expected to return same number of outputs but got:" + f"\n true branch returns {len(flat_true_outs)} item(s)" + f"\n false branch returns {len(flat_false_outs)} item(s)" + ) + + for i in range(0, len(flat_true_outs)): + true_out = flat_true_outs[i] + false_out = flat_false_outs[i] + + # Note that we need skip the check for requires_grad because we're after + # after autograd key during tracing, so the rquires_grad attribute of the tensors + # are no longer. See Note [invariants for node meta 'val'] + def _same_meta_except_requires_grad(true_out, false_out): + if true_out is None and false_out is None: + return True + elif true_out is None or false_out is None: + # Consider the following case: + # def true_fn(x, y): + # return x * y + # + # def false_fn(x, y): + # return x.sin() + # + # We'll get the following graphs for backward: + # def backward_true_fn(x, y, grad_out): + # return grad_out * y, grad_out * x + # + # def backward_false_fn(x, y, grad_out): + # retrun grad_out, None + # + # This suggests that when we make_fx into the backward graph, + # the output graph would produce outputs with metadata, this is undesirable. + # + # Ideally, we should provide an optional type to indicate that one of the branches might + # return None. But we'll just let it pass for now and let downstream/runtime handle. + # + # Note that this corner case should **only** happen when user want to trace backward graph because + # if it's foward, dynamo will error. + return True + true_meta = true_out.meta.get("tensor_meta", None) + false_meta = false_out.meta.get("tensor_meta", None) + return ( + true_meta.shape == false_meta.shape + and true_meta.dtype == false_meta.dtype + and true_meta.stride == false_meta.stride + ) + + if not _same_meta_except_requires_grad(true_out, false_out): + raise torch._dynamo.exc.CondOpArgsMismatchError( + f"Expected each tensor to have same metadata but got:" + f"\n {true_fn.__name__} returns {true_out.meta['tensor_meta']}" + f"\n {false_fn.__name__} returns {false_out.meta['tensor_meta']}" + ) + + i, true_name = unique_graph_id(proxy_mode, prefix="true_graph") + + false_name = f"false_graph_{i}" + assert not hasattr(proxy_mode.tracer.root, false_name) + + proxy_mode.tracer.root.register_module(true_name, true_graph) + proxy_mode.tracer.root.register_module(false_name, false_graph) + + args = (pred, true_graph, false_graph, operands) + + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, args) + + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", func_overload, proxy_args, {} + ) + + # At this point, we're *guaranteed* that whether an output came from the + # true or false branch is indistinguishable. So, as this is just for tracing + # purposes, choose the true branch. + + # TODO: the unbacked symbol allocations MUST NOT leak out, if you want to + # support this we need to arrange for the reenter_make_fx unbacked SymInts + # to be used, AND we need to arrange for some sort of unification between + # the two branches (but not really unification; e.g., if one branch + # returns [u0] and the other returns [5] this is OK but you MUST NOT + # conclude the result is 5. Also if one branch returns [3] and another + # branch returns [5] you can make it work by immediately allocating a new + # unbacked SymInt here). + ignore_fresh_unbacked = contextlib.nullcontext() + if (fake_mode := detect_fake_mode()) and fake_mode.shape_env: + ignore_fresh_unbacked = fake_mode.shape_env.ignore_fresh_unbacked_symbols() + + # TODO: Uhh.... it shouldn't matter, but changing this to true_fn results in + # a FakeTensorMode error : + # `Current active mode not registered` + # TODO Sometimes the operands are not completely FakeTensor, something seems went wrong in + # dynamo? Because of that it runs real computation sometimes and re-triggering downstream dispatch keys. + with ignore_fresh_unbacked: + out = false_fn(*operands) + + return track_tensor_tree(out, out_proxy, constant=None, tracer=proxy_mode.tracer) + + +@cond_op.py_impl(DispatchKey.CompositeExplicitAutograd) +def cond_op_dense(pred, true_fn, false_fn, operands): + assert all( + isinstance(o, (torch.Tensor, int)) for o in operands + ), f"Dense implementation operands must be a list of tensors and ints {operands}" + mode = _get_current_dispatch_mode() + assert mode is None, "Mode should never be enabled for CPU/CUDA key" + if pred: + return true_fn(*operands) + else: + return false_fn(*operands) + + +class CondAutogradOp(torch.autograd.Function): + @staticmethod + def forward( + ctx, + pred, + fw_true_graph, + fw_false_graph, + joint_true_graph, + joint_false_graph, + *operands, + ): + ctx._pred = pred + ctx._joint_true_graph = joint_true_graph + ctx._joint_false_graph = joint_false_graph + save_tensors_and_symints_for_backward(ctx, operands) + + with torch._C._AutoDispatchBelowAutograd(): + return cond_op(pred, fw_true_graph, fw_false_graph, operands) + + @staticmethod + def backward(ctx, *flat_grads): + operands = saved_tensors_and_symints(ctx) + + grads = cond_op( + ctx._pred, + ctx._joint_true_graph, + ctx._joint_false_graph, + flat_grads + operands, + ) + return None, None, None, None, None, *grads + + +@cond_op.py_impl(DispatchKey.Autograd) +def cond_autograd(pred, true_fn, false_fn, operands): + # A shortcut for the case where all inputs don't require gradient, + # we skip tracing the forward and backward graph. + if pytree.tree_all_only( + torch.Tensor, + lambda t: not t.requires_grad, # type: ignore[union-attr] + (pred, operands), + ): + with torch._C._AutoDispatchBelowAutograd(): + return cond_op(pred, true_fn, false_fn, operands) + + ( + fw_true_graph, + fw_false_graph, + joint_true_graph, + joint_false_graph, + ) = create_fw_bw_graph_branches(true_fn, false_fn, *operands) + flat_out = CondAutogradOp.apply( + pred, + fw_true_graph, + fw_false_graph, + joint_true_graph, + joint_false_graph, + *operands, + ) + return flat_out + + +@cond_op.py_impl(ProxyTorchDispatchMode) +def inner(mode, pred, true_fn, false_fn, operands): + return trace_cond(mode, cond_op, pred, true_fn, false_fn, operands) + + +@cond_op.py_impl(FakeTensorMode) +def cond_fake_tensor_mode(mode, pred, true_fn, false_fn, operands): + # Ignore here, because if you've gotten here but you're not manually + # tracing the inner graphs, that means that you intend to reuse the graph + # directly. Which means the old unbacked symbol bindings are appropriate. + # This strategy will not work if unbacked symbols can escape. + ignore_fresh_unbacked = contextlib.nullcontext() + if mode.shape_env: + ignore_fresh_unbacked = mode.shape_env.ignore_fresh_unbacked_symbols() + + with mode, ignore_fresh_unbacked: + true_outs = true_fn(*operands) + flat_true_outs = pytree.tree_leaves(true_outs) + flat_false_outs = pytree.tree_leaves(false_fn(*operands)) + if len(flat_true_outs) != len(flat_false_outs): + raise RuntimeError("Unmatched number of outputs from cond() branches.") + + for true_out, false_out in zip(flat_true_outs, flat_false_outs): + if true_out is None or false_out is None: + if true_out is None and false_out is None: + continue + raise torch._dynamo.exc.CondOpArgsMismatchError( + f"Expected both branches to return None:" + f"\n {true_fn.__name__} returns {true_out}" + f"\n {false_fn.__name__} returns {false_out}" + ) + true_meta = _extract_tensor_metadata(true_out) + false_meta = _extract_tensor_metadata(false_out) + if true_meta != false_meta: + raise torch._dynamo.exc.CondOpArgsMismatchError( + f"Expected each tensor to have same metadata but got:" + f"\n {true_fn.__name__} returns {true_meta}" + f"\n {false_fn.__name__} returns {false_meta}" + ) + return true_outs + + +@cond_op.py_functionalize_impl +def cond_func(ctx, pred, true_fn, false_fn, inputs): + unwrapped_inputs = ctx.unwrap_tensors(inputs) + unwrapped_pred = ctx.unwrap_tensors(pred) + with ctx.redispatch_to_next() as m: + functional_true = ctx.functionalize(_maybe_run_with_interpreter(true_fn)) + functional_false = ctx.functionalize(_maybe_run_with_interpreter(false_fn)) + pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch + for branch in [functional_true, functional_false]: + if _has_potential_branch_input_mutation( + branch, unwrapped_inputs, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException( + "One of torch.cond branch might be modifying the input! " + "Consider cloning the input before modifying it. " + ) + for branch in [true_fn, false_fn]: + if _has_potential_branch_input_alias( + branch, unwrapped_inputs, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException( + "One of torch.cond branch might be aliasing the input! " + "If you are returning a view of the input, please make sure " + "to clone it. " + ) + + cond_return = cond_op( + unwrapped_pred, functional_true, functional_false, unwrapped_inputs + ) + return ctx.wrap_tensors(cond_return) + + +@cond_op.py_impl(torch._C._functorch.TransformType.Vmap) +def cond_batch_rule(interpreter, pred, true_fn, false_fn, inputs): + assert isinstance( + inputs, (list, tuple) + ), "Cond inputs must be a list or tuple of tensors" + assert all( + isinstance(i, torch.Tensor) for i in inputs + ), "Cond inputs must be a list of tensors" + + pred_is_batched = isinstance(pred, torch.Tensor) and is_batchedtensor(pred) + pred_ = get_unwrapped(pred) if pred_is_batched else pred + + # unbatched tensors are not vmapped + tensors, in_dims = zip( + *[ + (get_unwrapped(t), maybe_get_bdim(t)) if is_batchedtensor(t) else (t, None) + for t in inputs + ] + ) + + if pred_is_batched: + # prepend "pred" and vmap everything + tensors = (pred_,) + tensors + in_dims = (0,) + in_dims + + def fn(p, *args): + t = true_fn(*args) + f = false_fn(*args) + return torch.where(p, t[0], f[0]) + + with interpreter.lower(): + result = torch.vmap(fn, in_dims=in_dims)(*tensors) + + else: + # predicate is known at this stage and it is a boolean expression or a + # tensor with one element. + true_fn = torch.vmap(true_fn, in_dims=in_dims) + false_fn = torch.vmap(false_fn, in_dims=in_dims) + + with interpreter.lower(): + result = cond_op(pred, true_fn, false_fn, tensors) + + if not isinstance(result, tuple): + result = (result,) + lvl = interpreter.level() + return tuple([_add_batch_dim(r, 0, lvl) for r in result]) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/effects.py b/lib/python3.12/site-packages/torch/_higher_order_ops/effects.py new file mode 100644 index 0000000000000000000000000000000000000000..3b05dbf84039367fb90e94a5e18a373efd27e4fc --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/effects.py @@ -0,0 +1,288 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +from enum import Enum +from typing import Any, Dict, Optional, Tuple, Union +from weakref import WeakKeyDictionary + +import torch +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._higher_order_ops.torchbind import call_torchbind +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ProxyTorchDispatchMode, + track_tensor_tree, +) + + +class _EffectType(Enum): + ORDERED = "Ordered" + + +OpType = Union[torch._ops.HigherOrderOperator, torch._ops.OpOverload] + + +SIDE_EFFECTS: "WeakKeyDictionary[OpType, _EffectType]" = WeakKeyDictionary( + { + torch.ops.aten._print.default: _EffectType.ORDERED, + call_torchbind: _EffectType.ORDERED, + } +) + + +def _register_effectful_op(op: OpType, effect: _EffectType): + assert isinstance( + op, (torch._ops.OpOverload, torch._ops.HigherOrderOperator) + ) and not has_aliasing(op) + if op in SIDE_EFFECTS and SIDE_EFFECTS[op] != effect: + raise RuntimeError( + f"Already registered effect type {SIDE_EFFECTS[op]} to op {op}, " + f"trying to register a different effect type {effect}." + ) + SIDE_EFFECTS[op] = effect + + +def _deregister_effectful_op(op: OpType): + if op not in SIDE_EFFECTS: + raise RuntimeError(f"Op {op} is not registered as effectful") + + del SIDE_EFFECTS[op] + + +class WithEffects(HigherOrderOperator): + """ + with_effects(token, op, args, kwargs) -> (new_token, op_results) + + This HOP helps ensure ordering between side effectful ops like prints or ops + using torchbind objects. This is needed to ensure a traced graph from + AOTAutograd is functional so that future optimization passes do not reorder + these operators. This is done through threading "effect tokens" through the + graph to enforce data dependence between side effectful ops. + + The tokens are basically dummy values (torch.tensor([])). We create a token + per "effect type", which are enumerated in the _EffectType enum. + """ + + def __init__(self) -> None: + super().__init__("with_effects") + + def __call__( + self, + token, + op: OpType, + *args: Tuple[Any, ...], + **kwargs: Dict[str, Any], + ) -> Tuple[Any, ...]: + assert isinstance(op, (torch._ops.HigherOrderOperator, torch._ops.OpOverload)) + assert not has_aliasing(op), "Ops with aliasing is not supported" + assert has_effects(op, args, kwargs) + assert isinstance(kwargs, dict) + return super().__call__(token, op, *args, **kwargs) + + +with_effects = WithEffects() + + +def has_aliasing(op: OpType): + # NOT FOR PUBLIC USE + if isinstance(op, torch._ops.HigherOrderOperator): + return op not in SIDE_EFFECTS + + for arg in op._schema.arguments: + if arg.alias_info is not None: + return True + for arg in op._schema.returns: + if arg.alias_info is not None: + return True + return False + + +def has_effects(op, args, kwargs) -> bool: + # Skip over the profiler's RecordFunction as they should not show up in the graph + _skip_ops = {torch.ops.profiler._record_function_exit._RecordFunction} + if op in _skip_ops: + return False + + return ( + isinstance(op, (torch._ops.HigherOrderOperator, torch._ops.OpOverload)) + and not has_aliasing(op) + and get_effect_key(op, args, kwargs) is not None + ) + + +def get_effect_key(op, args, kwargs) -> Optional[_EffectType]: + if op in SIDE_EFFECTS: + return SIDE_EFFECTS[op] + + for arg in args: + if isinstance(arg, torch.ScriptObject): + # Add it to the table so that next time we see the same op we don't + # have to parse through the args again + SIDE_EFFECTS[op] = _EffectType.ORDERED + return _EffectType.ORDERED + + return None + + +def new_token_tensor() -> torch.Tensor: + return torch.tensor([]) + + +@with_effects.py_impl(DispatchKey.CompositeExplicitAutograd) +def with_effects_dense( + token: torch.Tensor, + op: torch._ops.OpOverload, + *args: Tuple[Any, ...], + **kwargs: Dict[str, Any], +) -> Tuple[torch.Tensor, ...]: + out = op(*args, **kwargs) + new_token = new_token_tensor() + if isinstance(out, tuple): + return (new_token, *out) + return (new_token, out) + + +@with_effects.py_impl(FakeTensorMode) +def with_effects_fake( + mode, + token: torch.Tensor, + op: torch._ops.OpOverload, + *args: Tuple[Any, ...], + **kwargs: Dict[str, Any], +) -> Tuple[torch.Tensor, ...]: + with mode: + result = with_effects_dense(token, op, *args, **kwargs) + return result + + +@with_effects.py_impl(ProxyTorchDispatchMode) +def with_effects_proxy( + mode, + token: torch.Tensor, + op: torch._ops.OpOverload, + *args: Tuple[Any, ...], + **kwargs: Dict[str, Any], +) -> Tuple[torch.Tensor, ...]: + with disable_proxy_modes_tracing(): + out = with_effects(token, op, *args, **kwargs) + + proxy_token = mode.tracer.unwrap_proxy(token) + proxy_args = pytree.tree_map(mode.tracer.unwrap_proxy, args) + proxy_kwargs = pytree.tree_map(mode.tracer.unwrap_proxy, kwargs) + + from torch.fx.node import has_side_effect + + # To avoid the being DCEed by graph.eliminate_dead_code if they. + # don't have output or their outputs are not used. + has_side_effect(op) + + out_proxy = mode.tracer.create_proxy( + "call_function", + with_effects, + (proxy_token, op, *proxy_args), + proxy_kwargs, + ) + result = track_tensor_tree(out, out_proxy, constant=None, tracer=mode.tracer) + return result + + +with_effects.fallthrough(DispatchKey.AutogradCPU) +with_effects.fallthrough(DispatchKey.AutogradCUDA) + + +def _get_schema(op, args) -> torch.FunctionSchema: + if isinstance(op, torch._ops.OpOverload): + return op._schema + elif op == call_torchbind: + return getattr(args[0], args[1]).schema + else: + raise RuntimeError(f"Unable to get schema for op {op}") + + +def handle_effects( + allow_token_discovery: bool, + tokens: Dict[_EffectType, torch.Tensor], + op: OpType, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], +) -> Any: + """ + Args: + allow_token_discovery: Whether or not we are discovering tokens. If this + is true, we will create a token for every side effect type seen that + does not have a token assigned yet. If this is false, the tokens + should've all been created ahead of time, so we will error if there is + no token mapping to every effect type. + + tokens: Map of effect type to tokens. This is to chain operators of the + same effects together so that they do not get reordered in later + optimization passes. + """ + + # Get a token. We can't do `tokens.get(op, torch.tensor([]))` because + # this will create an empty tensor during proxy mode tracing if the token + # doesn't exist. But the tokens should always exist during proxy mode tracing. + key = get_effect_key(op, args, kwargs) + assert key is not None + if key not in tokens: + assert ( + allow_token_discovery + ), f"Could not find a token for effect {key} which came from the function {op}" + proxy_tensor_mode = torch._C._get_dispatch_mode( + torch._C._TorchDispatchModeKey.PROXY + ) + if proxy_tensor_mode is not None: + # If we discovered a new token during tracing, we are in backward. + # Then we patch the graph, adding additional tangents_token as input to the joint graph. + tracer = proxy_tensor_mode.tracer + + from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + track_tensor_tree, + ) + + with disable_proxy_modes_tracing(): + token_tensor = new_token_tensor() + + token_proxy = proxy_tensor_mode.tracer.create_proxy( + "placeholder", "tangents_token", (), {}, name="tangents_token" + ) + track_tensor_tree(token_tensor, token_proxy, constant=None, tracer=tracer) + + tokens[key] = token_tensor + else: + tokens[key] = new_token_tensor() + + token = tokens[key] + + from torch._subclasses.functional_tensor import PythonFunctionalizeAPI + + ctx = PythonFunctionalizeAPI() + + unwrapped_token = ctx.unwrap_tensors([token])[0] + unwrapped_args = ctx.unwrap_tensors(args) + unwrapped_kwargs = ctx.unwrap_tensors(kwargs) # type: ignore[arg-type] + with ctx.redispatch_to_next(): + (new_token, *unwrapped_outs) = with_effects( + unwrapped_token, op, *unwrapped_args, **unwrapped_kwargs + ) + + schema = _get_schema(op, unwrapped_args) + if len(schema.returns) == 0: + assert unwrapped_outs[0] is None + unwrapped_outs = None # type: ignore[assignment] + elif len(schema.returns) == 1: + assert len(unwrapped_outs) == 1 + unwrapped_outs = unwrapped_outs[0] + else: + assert len(unwrapped_outs) == len(schema.returns) + + # Add the newly created token into the tokens map for a following call to + # use this token. + wrapped_token = ctx.wrap_tensors(new_token) + assert isinstance(wrapped_token, torch.Tensor) + tokens[key] = wrapped_token + + return ctx.wrap_tensors(unwrapped_outs) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/executorch_call_delegate.py b/lib/python3.12/site-packages/torch/_higher_order_ops/executorch_call_delegate.py new file mode 100644 index 0000000000000000000000000000000000000000..a6ee5205ff4ebbde1240929727ec9e47b0a2433a --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/executorch_call_delegate.py @@ -0,0 +1,175 @@ +# mypy: allow-untyped-defs + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +from __future__ import annotations + +from typing import Any, cast + +import torch +import torch.utils._pytree as pytree +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + get_proxy_slot, + ProxyTorchDispatchMode, + track_tensor_tree, +) +from torch.utils._pytree import tree_flatten + + +class ExecutorchCallDelegate(HigherOrderOperator): + def __init__(self): + super().__init__("executorch_call_delegate") + + def __call__(self, lowered_module, *args): + return super().__call__(lowered_module, *args) + + +executorch_call_delegate = ExecutorchCallDelegate() +executorch_call_delegate.fallthrough(torch._C.DispatchKey.PythonDispatcher) +executorch_call_delegate.fallthrough(torch._C.DispatchKey.PythonTLSSnapshot) +executorch_call_delegate.fallthrough(torch._C.DispatchKey.ADInplaceOrView) +executorch_call_delegate.fallthrough(torch._C.DispatchKey.AutocastCPU) + +LOWERED_BACKEND_MODULE_TYPE = "LoweredBackendModule" + + +# pyre-ignore +def trace_call_delegate(proxy_mode, func_overload, lowered_module, *args): + # pyre-ignore + def _unwrap_proxy(e): + if not isinstance(e, (torch.Tensor, torch.SymInt, torch.SymFloat)): + return e + return get_proxy_slot( + cast(torch.Tensor, e), proxy_mode.tracer, e, lambda e: e.proxy # type: ignore[attr-defined] + ) + + if not is_lowered_module(lowered_module): + raise ValueError( + "executorch_call_delegate()'s first argument must be a LoweredBackendModule" + ) + + with disable_proxy_modes_tracing(): + out = call_delegate_cpu(lowered_module, *args) + + get_lowered_module_name(proxy_mode.tracer.root, lowered_module) + + node_args = (lowered_module, *args) + proxy_args = pytree.tree_map(_unwrap_proxy, node_args) + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", func_overload, proxy_args, {}, name="executorch_call_delegate" + ) + return track_tensor_tree(out, out_proxy, constant=None, tracer=proxy_mode.tracer) + + +@executorch_call_delegate.py_impl(torch._C.DispatchKey.CompositeExplicitAutograd) +# pyre-ignore +def call_delegate_cpu(lowered_module, *args): + # FX creates this immutable_dict/list concept. Get rid of this. + map_types = { + torch.fx.immutable_collections.immutable_dict: dict, + torch.fx.immutable_collections.immutable_list: list, + } + new_args = pytree.tree_map_only( + tuple(map_types.keys()), + lambda a: map_types[type(a)](a), + args, + lambda a: isinstance(a, tuple(map_types.keys())), + ) + return lowered_module.original_module.module()(*new_args) + + +@executorch_call_delegate.py_impl(torch._C.DispatchKey.Autograd) +# pyre-ignore +def call_delegate_autograd(lowered_module, *args): + # TODO: support autograd + flat_operands, _ = tree_flatten([lowered_module, *args]) + requires_grad = any( + f.requires_grad for f in flat_operands if isinstance(f, torch.Tensor) + ) + + with torch._C._ExcludeDispatchKeyGuard( + torch._C.DispatchKeySet(torch._C.DispatchKey.AutogradCPU) + ): + res = executorch_call_delegate(lowered_module, *args) + + if requires_grad: + # Create aliases of the output that has requires_grad=True. We need + # at least one of the inputs to err_fn to require grad so that the + # output will have a grad_fn. + + # pyre-ignore + def fake_requires_grad(var): + if var is not None: + var = var.detach() + if torch.is_floating_point(var) or torch.is_complex(var): + var.requires_grad = True + return var + + return pytree.tree_map_only(torch.Tensor, fake_requires_grad, res) + + return res + + +@executorch_call_delegate.py_impl(ProxyTorchDispatchMode) +# pyre-ignore +def call_delegate_proxy_torch_dispatch_mode(mode, lowered_module, *args): + res = trace_call_delegate(mode, executorch_call_delegate, lowered_module, *args) + return res + + +@executorch_call_delegate.py_impl(FakeTensorMode) +# pyre-ignore +def call_delegate_fake_tensor_mode(mode, lowered_module, *args): + with mode: + return call_delegate_cpu(lowered_module, *args) + + +@executorch_call_delegate.py_functionalize_impl +# pyre-ignore +def call_delegate_functionalize(ctx, lowered_module, *args): + unwrapped_args = tuple(ctx.unwrap_tensors(arg) for arg in args) + with ctx.redispatch_to_next(): + res = executorch_call_delegate(lowered_module, *unwrapped_args) + return ctx.wrap_tensors(res) + + +# pyre-ignore: Missing parameter annotation [2]: Parameter `obj` must have a type other than `Any`.Pyre +def is_lowered_module(obj: Any) -> bool: + """ + This function is added to avoid using isinstance(obj, + LoweredBackendModule) as it will import LoweredBackendModule, which may + cause a circular import. + """ + return type(obj).__name__ == LOWERED_BACKEND_MODULE_TYPE + + +def get_lowered_module_name( + root: torch.nn.Module, + # pyre-ignore: Undefined or invalid type [11]: Annotation `LoweredBackendModule` is not defined as a type. + lowered_module: LOWERED_BACKEND_MODULE_TYPE, # type: ignore[valid-type] +) -> str: + """ + Adds the given lowered_module into the given root module and returns the + name of the module added. + """ + # Find a qualifying name for the lowered submodule + qualname = None + i = 0 + while True: + qualname = f"lowered_module_{i}" + if not hasattr(root, qualname): + break + i += 1 + assert qualname is not None + + root.add_module(qualname, lowered_module) + return qualname diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/flex_attention.py b/lib/python3.12/site-packages/torch/_higher_order_ops/flex_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..e15e122fddb99767c397ca737daf7baf2b0bdf14 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/flex_attention.py @@ -0,0 +1,1149 @@ +import math +from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union + +import torch +import torch.utils._pytree as pytree +from torch import Tensor +from torch._C import DispatchKey +from torch._higher_order_ops.utils import ( + _has_potential_branch_input_mutation, + autograd_not_implemented, + reenter_make_fx, + save_tensors_and_symints_for_backward, + saved_tensors_and_symints, + UnsupportedAliasMutationException, + validate_subgraph_args_types, +) +from torch._ops import HigherOrderOperator +from torch._subclasses import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + make_fx, + ProxyTorchDispatchMode, + track_tensor_tree, +) +from torch.fx.graph_module import GraphModule + + +# Duplicate of _inductor/kernel/flex_attention.py to avoid circular import +def _construct_strides( + sizes: Sequence[int], + fill_order: Sequence[int], +) -> Sequence[int]: + """From a list of sizes and a fill order, construct the strides of the permuted tensor.""" + # Initialize strides + assert len(sizes) == len( + fill_order + ), "Length of sizes must match the length of the fill order" + strides = [0] * len(sizes) + + # Start with stride 1 for the innermost dimension + current_stride = 1 + + # Iterate through the fill order populating strides + for dim in fill_order: + strides[dim] = current_stride + current_stride *= sizes[dim] + + return strides + + +def _permute_strides(out: torch.Tensor, query_strides: Tuple[int, ...]) -> torch.Tensor: + """ + Create a new tensor with the same data and shape as the input, + but with strides permuted based on the input tensor's stride order. + + Args: + out (torch.Tensor): The output tensor of attention. + query_strides (List[int]): The stride order of the input query tensor + + Returns: + torch.Tensor: A new tensor with same shape and data as the input, + but with strides permuted based on the query tensor's stride order. + """ + from torch._inductor.ir import get_fill_order + + fill_order = get_fill_order(query_strides) + assert out.storage_offset() == 0, "Only support storage_offset == 0" + out_strides = _construct_strides(out.shape, fill_order) + new_out = out.new_empty(out.shape).as_strided(out.shape, out_strides) + new_out.copy_(out) + return new_out + + +class FlexAttentionHOP(HigherOrderOperator): + def __init__(self) -> None: + super().__init__("flex_attention", cacheable=True) + + def __call__( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + score_mod: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), + ) -> Tuple[torch.Tensor, torch.Tensor]: + validate_subgraph_args_types(score_mod_other_buffers + mask_mod_other_buffers) + return super().__call__( + query, + key, + value, + score_mod, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + + +flex_attention = FlexAttentionHOP() + + +class FlexAttentionBackwardHOP(HigherOrderOperator): + def __init__(self) -> None: + super().__init__("flex_attention_backward") + + def __call__( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + out: torch.Tensor, + logsumexp: torch.Tensor, + grad_out: torch.Tensor, + grad_logsumexp: torch.Tensor, + fw_graph: Union[Callable, GraphModule], + joint_graph: GraphModule, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), + ) -> Tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, Tuple[Optional[torch.Tensor], ...] + ]: + validate_subgraph_args_types(score_mod_other_buffers + mask_mod_other_buffers) + return super().__call__( + query, + key, + value, + out, + logsumexp, + grad_out, + grad_logsumexp, + fw_graph, + joint_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + + +flex_attention_backward = FlexAttentionBackwardHOP() + + +def _math_attention_inner( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + score_mod: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[torch.Tensor, torch.Tensor]: + from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex + + working_precision = torch.float64 if query.dtype == torch.float64 else torch.float32 + + scores = (query @ key.transpose(-2, -1)).to(dtype=working_precision) + + b = torch.arange(0, scores.size(0), device=scores.device) + h = torch.arange(0, scores.size(1), device=scores.device) + m = torch.arange(0, scores.size(2), device=scores.device) + n = torch.arange(0, scores.size(3), device=scores.device) + + captured_buffers_in_dim = (None,) * len(score_mod_other_buffers) + from torch.nn.attention.flex_attention import _vmap_for_bhqkv + + # first input is score + score_mod = _vmap_for_bhqkv(score_mod, prefix=(0,), suffix=captured_buffers_in_dim) + + mask_mod = block_mask[-1] + mask_mod_in_dim_buffers = (None,) * len(mask_mod_other_buffers) + mask_mod = _vmap_for_bhqkv(mask_mod, prefix=(), suffix=mask_mod_in_dim_buffers) + + with TransformGetItemToIndex(): + scores = (scores * scale).to(working_precision) + post_mod_scores = torch.where( + mask_mod(b, h, m, n, *mask_mod_other_buffers), + score_mod(scores, b, h, m, n, *score_mod_other_buffers), + torch.tensor(-float("inf"), dtype=working_precision, device=scores.device), + ) + + return scores, post_mod_scores + + +def math_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + score_mod: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[torch.Tensor, torch.Tensor]: + """Eager implementation + + This implementation uses vmap to vectorize the score_mod function over the batch, head, m, and n dimensions. + We then apply the vectorized score_mod function to the scores matrix. Each wrap of vmap applies one of the + batch, head, m, or n dimensions. We need to apply vmap 4 times to vectorized over all 4 dimensions. + + Args: + query: The query tensor + key: The key tensor + value: The value tensor + score_mod: The score_mod function + other_buffers: Other buffers that are passed to the score_mod function + """ + # broadcast query & key along head dim for GQA + G = query.size(1) // key.size(1) + value = torch.repeat_interleave(value, G, dim=1) + key = torch.repeat_interleave(key, G, dim=1) + + Bq, Bkv = query.size(0), key.size(0) + if not ((Bq == Bkv) or (Bq > 1 and Bkv == 1)): + raise RuntimeError(f"Bq and Bkv must broadcast. Got Bq={Bq} and Bkv={Bkv}") + + key = key.expand((Bq, *key.size()[1:])) + value = value.expand((Bq, *value.size()[1:])) + + _, post_mod_scores = _math_attention_inner( + query, + key, + value, + score_mod, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + + # Set fully masked rows' sumexp to 0.0 + logsumexp = post_mod_scores.logsumexp(dim=-1) + masked_rows = torch.all(post_mod_scores == -float("inf"), dim=-1) + logsumexp = torch.where(masked_rows, -float("inf"), logsumexp) + + post_mod_scores = torch._safe_softmax(post_mod_scores, dim=-1) + + return post_mod_scores.to(query.dtype) @ value, logsumexp / math.log(2) + + +@flex_attention.py_impl(DispatchKey.CompositeExplicitAutograd) +def sdpa_dense( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + score_mod: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[torch.Tensor, torch.Tensor]: + out, lse = math_attention( + query, + key, + value, + score_mod, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + out = _permute_strides(out, query.stride()) + return out, lse + + +def trace_flex_attention( + proxy_mode: ProxyTorchDispatchMode, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + score_mod: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[torch.Tensor, torch.Tensor]: + """Traces the flex_attention operator with the given score_mod function and other_buffers. + + Trace SDPA will call make_fx with "fake" example vals and then trace the score_mod function + This will produce a GraphModule that will be stored on the root tracer as "sdpa_score". We + access this graph module in inductor to inline the score_mod function to the triton template. + """ + from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex + + example_out = flex_attention( + query, + key, + value, + score_mod, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + example_vals = [query.new_zeros((), requires_grad=query.requires_grad)] + [ + query.new_zeros((), dtype=torch.int) for _ in range(4) + ] + mask_example_vals = [query.new_zeros((), dtype=torch.int) for _ in range(4)] + mask_mod = block_mask[-1] + with TransformGetItemToIndex(): + score_graph = reenter_make_fx(score_mod)( + *example_vals, *score_mod_other_buffers + ) + mask_graph = reenter_make_fx(mask_mod)( + *mask_example_vals, *mask_mod_other_buffers + ) + assert isinstance(proxy_mode.tracer, torch.fx.Tracer) + block_mask = block_mask[:-1] + (mask_graph,) + qualname = proxy_mode.tracer.get_fresh_qualname("sdpa_score") + proxy_mode.tracer.root.register_module(qualname, score_graph) + mask_qualname = proxy_mode.tracer.get_fresh_qualname("sdpa_mask") + proxy_mode.tracer.root.register_module(mask_qualname, mask_graph) + node_args = ( + query, + key, + value, + score_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, node_args) + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", flex_attention, proxy_args, {} + ) + return track_tensor_tree( + example_out, out_proxy, constant=None, tracer=proxy_mode.tracer + ) + + +@flex_attention.py_impl(ProxyTorchDispatchMode) +def flex_attention_proxy_torch_dispatch_mode( + mode: ProxyTorchDispatchMode, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + score_mod: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[torch.Tensor, torch.Tensor]: + assert mode is not None, "Mode should always be enabled for python fallback key" + return trace_flex_attention( + mode, + query, + key, + value, + score_mod, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + + +@flex_attention.py_functionalize_impl +def flex_attention_functionalize( + ctx: torch._subclasses.functional_tensor.BaseFunctionalizeAPI, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + score_mod: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[torch.Tensor, torch.Tensor]: + """Defines the functionalization rules for the flex_attention operator. + + Write now we are unwrapping each tensor and then redispatching to the next, however we want to + guard against any mutations in the score_mod function, to the other_buffers since those + are free variables. + """ + from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex + + query_unwrapped = ctx.unwrap_tensors(query) + key_unwrapped = ctx.unwrap_tensors(key) + value_unwrapped = ctx.unwrap_tensors(value) + block_mask_unwrapped = ctx.unwrap_tensors(block_mask) + score_mod_other_buffers_unwrapped = ctx.unwrap_tensors(score_mod_other_buffers) + mask_mod_other_buffers_unwrapped = ctx.unwrap_tensors(mask_mod_other_buffers) + + # Appease the mypy overlords + assert isinstance(query_unwrapped, torch.Tensor) + assert isinstance(key_unwrapped, torch.Tensor) + assert isinstance(value_unwrapped, torch.Tensor) + assert isinstance(block_mask_unwrapped, tuple) + assert isinstance(score_mod_other_buffers_unwrapped, tuple) + assert isinstance(mask_mod_other_buffers_unwrapped, tuple) + + example_vals = ( + [query_unwrapped.new_zeros(())] + + [query_unwrapped.new_zeros((), dtype=torch.int) for _ in range(4)] + + list(score_mod_other_buffers_unwrapped) + ) + with ctx.redispatch_to_next() as m: + functional_score_mod = ctx.functionalize(score_mod) + pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch + with TransformGetItemToIndex(): + mutates = _has_potential_branch_input_mutation( + functional_score_mod, example_vals, pre_dispatch + ) + # The only care about mutations of existing buffers since we can't replay these. + # However, we can just error if anything is detected + if mutates: + raise UnsupportedAliasMutationException("Mutations detected in score_mod") + + out = flex_attention( + query_unwrapped, + key_unwrapped, + value_unwrapped, + functional_score_mod, + block_mask_unwrapped, + scale, + kernel_options, + score_mod_other_buffers_unwrapped, + mask_mod_other_buffers_unwrapped, + ) + return ctx.wrap_tensors(out) # type: ignore[return-value, arg-type] + + +@flex_attention.py_impl(FakeTensorMode) +def flex_attention_fake_tensor_mode( + mode: FakeTensorMode, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + score_mod: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[torch.Tensor, torch.Tensor]: + with mode: + v_head_dim = value.size(-1) + batch_size, num_heads, seq_len_q, q_head_dim = query.shape + logsumexp = query.new_empty( + batch_size, num_heads, seq_len_q, dtype=torch.float32 + ) + out_shape = (batch_size, num_heads, seq_len_q, v_head_dim) + out = query.new_empty(out_shape) + out = _permute_strides(out, query.stride()) + return out, logsumexp + + +# ---------------------------- Autograd Implementation ---------------------------- +def create_fw_bw_graph( + score_mod: Callable, + index_values: Tuple[Tensor, Tensor, Tensor, Tensor, Tensor], + other_buffers: Tuple[Tensor, ...], +) -> Tuple[Callable, Callable]: + # See Note:[HOP create fw_bw graph] + + # All of these imports need to be here in order to avoid circular dependencies + from torch._dispatch.python import suspend_functionalization + from torch._functorch.aot_autograd import AOTConfig, create_joint + from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode + from torch._subclasses.functional_tensor import disable_functional_mode + from torch.fx.experimental.proxy_tensor import disable_proxy_modes_tracing + + dummy_aot_config = AOTConfig( + fw_compiler=None, # type: ignore[arg-type] + bw_compiler=None, # type: ignore[arg-type] + partition_fn=None, # type: ignore[arg-type] + decompositions={}, + num_params_buffers=0, + aot_id=0, + keep_inference_input_mutations=False, + ) + + with suspend_functionalization(), disable_functional_mode(): + with disable_proxy_modes_tracing(): + + def _from_fun( + t: Union[Tensor, torch.SymInt, int] + ) -> Union[Tensor, torch.SymInt, int]: + if isinstance(t, torch.Tensor): + return torch.empty_strided( + t.size(), + t.stride(), + device=t.device, + dtype=t.dtype, + requires_grad=t.requires_grad, + ) + return t + + # If someone runs this hop under the default compiler backend ("eager") + # Then this path will be run with the actual user inputs. We convert them + # to fake tensors in order to not perform any actual compute. + from torch._guards import detect_fake_mode + + fake_mode = detect_fake_mode(index_values) + if fake_mode is None: + fake_mode = FakeTensorMode(allow_non_fake_inputs=True) + + with fake_mode: + unwrapped_score_mod_indexes = pytree.tree_map(_from_fun, index_values) + unwrapped_other_buffers = pytree.tree_map(_from_fun, other_buffers) + + assert all( + isinstance(t, (FakeTensor, int, torch.SymInt)) + for t in unwrapped_score_mod_indexes + unwrapped_other_buffers + ) + + example_flat_out = pytree.tree_map( + _from_fun, + score_mod(*unwrapped_score_mod_indexes, *unwrapped_other_buffers), + ) + if not isinstance(example_flat_out, torch.Tensor): + raise RuntimeError( + "Expected output of score_mod to be a tensor." + f"Got type {type(example_flat_out)}." + ) + example_grad = _from_fun(example_flat_out) + + def joint_f( + score: Tensor, + b: Tensor, + h: Tensor, + m: Tensor, + n: Tensor, + example_grad: Tensor, + *other_buffers: Tuple[Tensor, ...], + ) -> Tuple[Tensor, ...]: + def fw_with_masks( + *args: Tuple[Tensor, ...] + ) -> Tuple[Tuple[Tensor], Tuple[bool]]: + fw_out = score_mod(*args) + out_requires_grad = fw_out.requires_grad + return ((fw_out,), (out_requires_grad,)) + + joint = create_joint(fw_with_masks, aot_config=dummy_aot_config) + args = [score, b, h, m, n] + list(other_buffers) + optional_grad = [example_grad] if example_grad.requires_grad else [] + _, grads = joint(args, optional_grad) + + return grads + + joint_graph = make_fx(joint_f)( + *unwrapped_score_mod_indexes, example_grad, *unwrapped_other_buffers + ) + return score_mod, joint_graph + + +class FlexAttentionAutogradOp(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + query: Tensor, + key: Tensor, + value: Tensor, + fw_graph: Callable, + joint_graph: Callable, + block_mask: Tuple[Any, ...], + scale: float, + kernel_options: Dict[str, Any], + mask_mod_other_buffers: Tuple[Any, ...], + *score_mod_other_buffers: Tuple[Any, ...], + ) -> Tuple[torch.Tensor, torch.Tensor]: + any_buffer_requires_grad = any( + buffer.requires_grad + for buffer in mask_mod_other_buffers + if isinstance(buffer, torch.Tensor) + ) + assert ( + not any_buffer_requires_grad + ), "Captured buffers from mask mod that require grad are not supported." + ctx._fw_graph = fw_graph + ctx._joint_graph = joint_graph + ctx._mask_graph = block_mask[-1] + ctx.scale = scale + ctx.kernel_options = kernel_options + ctx._score_mod_other_buffers_len = len(score_mod_other_buffers) + with torch._C._AutoDispatchBelowAutograd(): + out, logsumexp = flex_attention( + query, + key, + value, + fw_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + + save_tensors_and_symints_for_backward( + ctx, + ( + query, + key, + value, + out, + logsumexp, + *block_mask[:-1], + *score_mod_other_buffers, + *mask_mod_other_buffers, + ), + ) + return out, logsumexp + + @staticmethod + def backward(ctx: Any, grad_out: Tensor, grad_logsumexp: Tensor) -> Tuple[Optional[Tensor], ...]: # type: ignore[override] + fw_args = saved_tensors_and_symints(ctx) + ( + query, + key, + value, + out, + logsumexp, + query_lengths, + kv_lengths, + kv_num_blocks, + kv_indices, + full_kv_num_blocks, + full_kv_indices, + q_num_blocks, + q_indices, + full_q_num_blocks, + full_q_indices, + Q_BLOCK_SIZE, + KV_BLOCK_SIZE, + *other_buffers, + ) = fw_args + fw_graph = ctx._fw_graph + joint_graph = ctx._joint_graph + mask_graph = ctx._mask_graph + scale = ctx.scale + kernel_options = ctx.kernel_options + score_mod_other_buffers = tuple( + other_buffers[: ctx._score_mod_other_buffers_len] + ) + mask_mod_other_buffers = tuple( + other_buffers[ctx._score_mod_other_buffers_len :] + ) + # We have asserted that mask_mod_other_buffers do not require grad, + # but score_mod_other_buffers can require grad. + none_grads = [None] * 6 + ( + grad_query, + grad_key, + grad_value, + grad_score_mod_captured, + ) = flex_attention_backward( + query, + key, + value, + out, + logsumexp, + grad_out, + grad_logsumexp, + fw_graph, + joint_graph, + ( + query_lengths, + kv_lengths, + kv_num_blocks, + kv_indices, + full_kv_num_blocks, + full_kv_indices, + q_num_blocks, + q_indices, + full_q_num_blocks, + full_q_indices, + Q_BLOCK_SIZE, + KV_BLOCK_SIZE, + mask_graph, + ), + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + return grad_query, grad_key, grad_value, *none_grads, *grad_score_mod_captured + + +@flex_attention.py_impl(DispatchKey.Autograd) +def flex_attention_autograd( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + score_mod: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple[Tensor, ...] = (), + mask_mod_other_buffers: Tuple[Tensor, ...] = (), +) -> Tuple[torch.Tensor, torch.Tensor]: + from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex + + with TransformGetItemToIndex(): + input_requires_grad = any( + isinstance(t, torch.Tensor) and t.requires_grad + for t in (query, key, value, *score_mod_other_buffers) + ) + if torch.is_grad_enabled() and input_requires_grad: + example_vals = ( + query.new_zeros((), requires_grad=input_requires_grad), + query.new_zeros((), dtype=torch.int), + query.new_zeros((), dtype=torch.int), + query.new_zeros((), dtype=torch.int), + query.new_zeros((), dtype=torch.int), + ) + fw_graph, bw_graph = create_fw_bw_graph( + score_mod, example_vals, score_mod_other_buffers + ) + else: + fw_graph, bw_graph = score_mod, None + out, logsumexp = FlexAttentionAutogradOp.apply( + query, + key, + value, + fw_graph, + bw_graph, + block_mask, + scale, + kernel_options, + mask_mod_other_buffers, + *score_mod_other_buffers, + ) + return out, logsumexp + + +# ---------------------------- Backward HOP Implementation ---------------------------- + + +@flex_attention_backward.py_impl(DispatchKey.CompositeExplicitAutograd) +def sdpa_dense_backward( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + out: torch.Tensor, + logsumexp: torch.Tensor, + grad_out: torch.Tensor, + grad_logsumexp: torch.Tensor, + fw_graph: Callable, # GraphModule type hint? + joint_graph: Callable, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple, + mask_mod_other_buffers: Tuple, +) -> Tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, Tuple[Optional[torch.Tensor], ...] +]: + from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex + + # Get outputs before calling repeat interleave + actual_grad_query = torch.empty_like(query) + actual_grad_key = torch.empty_like(key) + actual_grad_value = torch.empty_like(value) + + def _maybe_new_buffer( + buffer: Union[torch.Tensor, torch.SymInt, int] + ) -> Optional[Union[torch.Tensor, torch.SymInt, int]]: + if isinstance(buffer, torch.Tensor): + return torch.empty_like(buffer) if buffer.requires_grad else None + return buffer + + actual_grad_score_mod_captured = [ + _maybe_new_buffer(buffer) for buffer in score_mod_other_buffers + ] + + Bq, Bkv = query.size(0), key.size(0) + if not ((Bq == Bkv) or (Bq > 1 and Bkv == 1)): + raise RuntimeError(f"Bq and Bkv must broadcast. Got Bq={Bq} and Bkv={Bkv}") + + key = key.expand((Bq, *key.size()[1:])) + value = value.expand((Bq, *value.size()[1:])) + + G = query.size(1) // key.size(1) + key = torch.repeat_interleave(key, G, dim=1) + value = torch.repeat_interleave(value, G, dim=1) + + # We're undoing the log -> log2 change of base in the forwards + logsumexp = logsumexp * math.log(2) + # The backwards formula for the log -> log2 change of base in the forwards + grad_logsumexp = grad_logsumexp / math.log(2) + scores, post_mod_scores = _math_attention_inner( + query, + key, + value, + fw_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + masked_out_rows = logsumexp == -float("inf") + softmax_scores = torch.exp(post_mod_scores - logsumexp.unsqueeze(-1)) + softmax_scores = torch.where(masked_out_rows.unsqueeze(-1), 0, softmax_scores) + + grad_value = softmax_scores.to(query.dtype).transpose(-2, -1) @ grad_out + + grad_softmax_scores = grad_out @ value.transpose(-2, -1) + + sum_scores = torch.sum(out * grad_out, -1, keepdim=True) + grad_score_mod = softmax_scores * ( + grad_softmax_scores - sum_scores + grad_logsumexp.unsqueeze(-1) + ) + + b = torch.arange(0, scores.size(0), device=scores.device) + h = torch.arange(0, scores.size(1), device=scores.device) + m = torch.arange(0, scores.size(2), device=scores.device) + n = torch.arange(0, scores.size(3), device=scores.device) + + mask_graph = block_mask[-1] + # Gradient of the inline score_mod function, with respect to the scores + captured_buffers_in_dim = (None,) * len(score_mod_other_buffers) + out_dims = [0, None, None, None, None] + [None] * len(score_mod_other_buffers) + from torch.nn.attention.flex_attention import _vmap_for_bhqkv + + # inputs are [score, b, h, q_idx, kv_idx, gradOut, ...] + # score and gradOut are "fully" batched + joint_score_mod = _vmap_for_bhqkv( + joint_graph, + prefix=(0,), + suffix=(0,) + captured_buffers_in_dim, + out_dims=out_dims, + ) + with TransformGetItemToIndex(): + grad_scores, _, _, _, _, *grad_score_mod_captured = joint_score_mod( + scores, b, h, m, n, grad_score_mod, *score_mod_other_buffers + ) + grad_scores = grad_scores * scale + grad_scores = grad_scores.to(query.dtype) + + mask_mod = _vmap_for_bhqkv( + mask_graph, prefix=(), suffix=(None,) * len(mask_mod_other_buffers) + ) + with TransformGetItemToIndex(): + mask_scores = mask_mod(b, h, m, n, *mask_mod_other_buffers) + grad_scores = torch.where( + mask_scores, grad_scores, torch.tensor(0, dtype=query.dtype) + ) + + grad_query = grad_scores @ key + grad_key = grad_scores.transpose(-2, -1) @ query + + # Reduce DK, DV along broadcasted heads. + grad_key = grad_key.view( + grad_key.size(0), -1, G, grad_key.size(-2), grad_key.size(-1) + ) + grad_value = grad_value.view( + grad_value.size(0), -1, G, grad_value.size(-2), grad_value.size(-1) + ) + + grad_key = torch.sum(grad_key, 2, keepdim=False) + grad_value = torch.sum(grad_value, 2, keepdim=False) + + if Bq != Bkv: + assert ( + Bq > 1 and Bkv == 1 + ), f"Bq and Bkv must broadcast. Got Bq={Bq} and Bkv={Bkv}" + + # Reduce DK, DV along broadcasted batches. + grad_key = torch.sum(grad_key, 0, keepdim=True) + grad_value = torch.sum(grad_value, 0, keepdim=True) + + actual_grad_query.copy_(grad_query) + actual_grad_key.copy_(grad_key) + actual_grad_value.copy_(grad_value) + score_mod_other_buffer_grads = [ + actual_grad.copy_(grad) if isinstance(actual_grad, torch.Tensor) else None + for actual_grad, grad in zip( + actual_grad_score_mod_captured, grad_score_mod_captured + ) + ] + + return ( + actual_grad_query, + actual_grad_key, + actual_grad_value, + tuple(score_mod_other_buffer_grads), + ) + + +def trace_flex_attention_backward( + proxy_mode: ProxyTorchDispatchMode, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + out: torch.Tensor, + logsumexp: torch.Tensor, + grad_out: torch.Tensor, + grad_logsumexp: torch.Tensor, + fw_graph: Union[Callable, GraphModule], + joint_graph: GraphModule, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, Tuple[Optional[torch.Tensor], ...] +]: + """We already have the forward graph and joint graph from the forward pass, so we create a proxy attach both graphs""" + from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex + + example_out = flex_attention_backward( + query, + key, + value, + out, + logsumexp, + grad_out, + grad_logsumexp, + fw_graph, + joint_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + + requires_grad = any(pytree.tree_map(lambda x: x.requires_grad, (query, key))) + fw_example_vals = [query.new_zeros((), requires_grad=requires_grad)] + [ + query.new_zeros((), dtype=torch.int) for _ in range(4) + ] + bw_example_vals = fw_example_vals + [query.new_zeros(())] + mask_example_vals = [query.new_zeros((), dtype=torch.int) for _ in range(4)] + mask_graph = block_mask[-1] + with TransformGetItemToIndex(): + fw_graph = reenter_make_fx(fw_graph)(*fw_example_vals, *score_mod_other_buffers) + joint_graph = reenter_make_fx(joint_graph)( + *bw_example_vals, *score_mod_other_buffers + ) + mask_graph = reenter_make_fx(mask_graph)( + *mask_example_vals, *mask_mod_other_buffers + ) + assert isinstance(proxy_mode.tracer, torch.fx.Tracer) + block_mask = block_mask[:-1] + (mask_graph,) + + qualname = proxy_mode.tracer.get_fresh_qualname("fw_graph") + proxy_mode.tracer.root.register_module(qualname, fw_graph) # type: ignore[arg-type] + qualname = proxy_mode.tracer.get_fresh_qualname("joint_graph") + proxy_mode.tracer.root.register_module(qualname, joint_graph) + qualname = proxy_mode.tracer.get_fresh_qualname("mask_graph") + proxy_mode.tracer.root.register_module(qualname, mask_graph) + + node_args = ( + query, + key, + value, + out, + logsumexp, + grad_out, + grad_logsumexp, + fw_graph, + joint_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, node_args) + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", + flex_attention_backward, + proxy_args, + {}, + name="flex_attention_backward", + ) + return track_tensor_tree( + example_out, out_proxy, constant=None, tracer=proxy_mode.tracer + ) + + +@flex_attention_backward.py_impl(ProxyTorchDispatchMode) +def flex_attention_backward_proxy_torch_dispatch_mode( + mode: ProxyTorchDispatchMode, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + out: torch.Tensor, + logsumexp: torch.Tensor, + grad_out: torch.Tensor, + grad_logsumexp: torch.Tensor, + fw_graph: Union[Callable, GraphModule], + joint_graph: GraphModule, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, Tuple[Optional[torch.Tensor], ...] +]: + assert mode is not None, "Mode should always be enabled for python fallback key" + return trace_flex_attention_backward( + mode, + query, + key, + value, + out, + logsumexp, + grad_out, + grad_logsumexp, + fw_graph, + joint_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + + +@flex_attention_backward.py_functionalize_impl +def flex_attention_backward_functionalize( + ctx: torch._subclasses.functional_tensor.BaseFunctionalizeAPI, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + out: torch.Tensor, + logsumexp: torch.Tensor, + grad_out: torch.Tensor, + grad_logsumexp: torch.Tensor, + fw_graph: Union[Callable, GraphModule], + joint_graph: GraphModule, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, Tuple[Optional[torch.Tensor], ...] +]: + """Defines the functionalization rules for the flex_attention operator. + + Write now we are unwrapping each tensor and then redispatching to the next, + since we know that the forward score mod function is assured to be free of mutations + to the other_buffers, we skip that mutate check and go straight to redispatching. + """ + query_unwrapped = ctx.unwrap_tensors(query) + key_unwrapped = ctx.unwrap_tensors(key) + value_unwrapped = ctx.unwrap_tensors(value) + out_unwrapped = ctx.unwrap_tensors(out) + logsumexp_unwrapped = ctx.unwrap_tensors(logsumexp) + grad_out_unwrapped = ctx.unwrap_tensors(grad_out) + grad_logsumexp_unwrapped = ctx.unwrap_tensors(grad_logsumexp) + block_mask_unwrapped = ctx.unwrap_tensors(block_mask) + score_mod_other_buffers_unwrapped = ctx.unwrap_tensors(score_mod_other_buffers) + mask_mod_other_buffers_unwrapped = ctx.unwrap_tensors(mask_mod_other_buffers) + + # Appease the mypy overlords + assert isinstance(query_unwrapped, torch.Tensor) + assert isinstance(key_unwrapped, torch.Tensor) + assert isinstance(value_unwrapped, torch.Tensor) + assert isinstance(out_unwrapped, torch.Tensor) + assert isinstance(logsumexp_unwrapped, torch.Tensor) + assert isinstance(grad_out_unwrapped, torch.Tensor) + assert isinstance(grad_logsumexp_unwrapped, torch.Tensor) + assert isinstance(block_mask_unwrapped, tuple) + assert isinstance(score_mod_other_buffers_unwrapped, tuple) + assert isinstance(mask_mod_other_buffers_unwrapped, tuple) + + with ctx.redispatch_to_next() as m: + functional_fw_graph = ctx.functionalize(fw_graph) + functional_joint_graph = ctx.functionalize(joint_graph) + + ( + grad_query, + grad_key, + grad_value, + grad_score_mod_captured, + ) = flex_attention_backward( + query_unwrapped, + key_unwrapped, + value_unwrapped, + out_unwrapped, + logsumexp_unwrapped, + grad_out_unwrapped, + grad_logsumexp_unwrapped, + functional_fw_graph, # type: ignore[arg-type] + functional_joint_graph, # type: ignore[arg-type] + block_mask_unwrapped, + scale, + kernel_options, + score_mod_other_buffers_unwrapped, + mask_mod_other_buffers_unwrapped, + ) + + return ctx.wrap_tensors((grad_query, grad_key, grad_value, grad_score_mod_captured)) # type: ignore[return-value,arg-type] + + +@flex_attention_backward.py_impl(FakeTensorMode) +def flex_attention_backward_fake_tensor_mode( + mode: FakeTensorMode, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + out: torch.Tensor, + logsumexp: torch.Tensor, + grad_out: torch.Tensor, + grad_logsumexp: torch.Tensor, + fw_graph: Union[Callable, GraphModule], + joint_graph: GraphModule, + block_mask: Tuple, + scale: float, + kernel_options: Dict[str, Any], + score_mod_other_buffers: Tuple = (), + mask_mod_other_buffers: Tuple = (), +) -> Tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, Tuple[Optional[torch.Tensor], ...] +]: + with mode: + grad_query = torch.empty_like(query) + grad_key = torch.empty_like(key) + grad_value = torch.empty_like(value) + grad_score_mod_captured = tuple( + [ + torch.empty_like(buffer) + if isinstance(buffer, torch.Tensor) and buffer.requires_grad + else None + for buffer in score_mod_other_buffers + ] + ) + return grad_query, grad_key, grad_value, grad_score_mod_captured + + +flex_attention_backward.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(flex_attention_backward, deferred_error=True) +) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/hints_wrap.py b/lib/python3.12/site-packages/torch/_higher_order_ops/hints_wrap.py new file mode 100644 index 0000000000000000000000000000000000000000..c211d405614636085575369e4a3738d3ced3fe02 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/hints_wrap.py @@ -0,0 +1,151 @@ +# mypy: allow-untyped-defs +import torch +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._higher_order_ops.utils import ( + _has_potential_branch_input_alias, + _has_potential_branch_input_mutation, + autograd_not_implemented, + reenter_make_fx, + unique_graph_id, + UnsupportedAliasMutationException, +) +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode, track_tensor_tree + + +# used for wrapping a function/op with context hints +class HintsWrapper(HigherOrderOperator): + def __init__(self): + super().__init__("hints_wrapper") + + def __call__(self, body_fn, args, kwargs, hints): + r""" + Call implementation of hints_wrapper + + Args: + body_fn (Callable): A callable function that is within the scope + that is being traced. + + args (Tuple of torch.Tensor/int/float/bool): A tuple of inputs to + body_fn. + + kwargs (dict): Keyword argument to the body_fn. + + hints (dict): A dict of context hints which could be passed to + backend compiler. + """ + if not isinstance(args, tuple): + raise RuntimeError(f"args must be a tuple, got {type(args)}") + + if not all(isinstance(t, (torch.Tensor, int, float, bool)) for t in args): + raise RuntimeError( + "args must be a tuple of tensors, ints, floats, or bools, got " + f"{args}" + ) + + if not isinstance(kwargs, dict): + raise RuntimeError(f"kwargs must be a dict, got {type(kwargs)}") + + if len(kwargs) > 0: + raise RuntimeError( + f"kwargs except for hints are not supported, got {kwargs}" + ) + + if not isinstance(hints, dict): + raise RuntimeError(f"hints must be a dict, got {type(hints)}") + + for k, v in hints.items(): + if not isinstance(k, str): + raise RuntimeError(f"hints key must be a str, got {k}.") + + if not isinstance(v, (int, float, bool, str)): + raise RuntimeError( + "hints must be a dict containing int, float, bool or str " + f"value, got value {v} for key {k}." + ) + + return super().__call__(body_fn, args, kwargs, hints) + + +hints_wrapper = HintsWrapper() + + +@hints_wrapper.py_impl(DispatchKey.CompositeExplicitAutograd) +def hints_wrapper_dense(body_fn, args, kwargs, hints): + return body_fn(*args, **kwargs) + + +hints_wrapper.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(hints_wrapper, deferred_error=True) +) + + +@hints_wrapper.py_impl(FakeTensorMode) +def hints_wrapper_fake_tensor_mode(mode, body_func, args, kwargs, hints): + flat_args = pytree.tree_leaves(args) + with mode: + return body_func(*flat_args, **kwargs) + + +@hints_wrapper.py_functionalize_impl +def hints_wrapper_functionalize(ctx, body_fn, args, kwargs, hints): + unwrapped_args = ctx.unwrap_tensors(args) + unwrapped_kwargs = ctx.unwrap_tensors(kwargs) + unwrapped_hints = ctx.unwrap_tensors(hints) + with ctx.redispatch_to_next(): + functional_body_fn = ctx.functionalize(body_fn) + pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch + if _has_potential_branch_input_mutation( + functional_body_fn, unwrapped_args, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException( + "body_fn of hints_wrapper might be modifying the input!" + ) + if _has_potential_branch_input_alias( + functional_body_fn, unwrapped_args, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException( + "body_fn of hints_wrapper might be aliasing the input!" + ) + outputs = hints_wrapper( + functional_body_fn, + unwrapped_args, + unwrapped_kwargs, + unwrapped_hints, + ) + return ctx.wrap_tensors(outputs) + + +def trace_hints_wrapper(proxy_mode, hints_wrapper, body_fn, args, kwargs, hints): + flat_args = tuple(pytree.tree_leaves(args)) + body_graph = reenter_make_fx(body_fn)(*flat_args, **kwargs) + + _, body_graph_name = unique_graph_id(proxy_mode, prefix="hints_wrapper_body_graph") + proxy_mode.tracer.root.register_module(body_graph_name, body_graph) + + new_args: tuple = (body_graph, flat_args, {}) + # merge hints into kwargs + new_kwargs = {} + new_kwargs["hints"] = hints + + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, new_args) + proxy_kwargs = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, new_kwargs) + + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", hints_wrapper, proxy_args, proxy_kwargs, name="hints_wrapper" + ) + + out = body_fn(*flat_args, **kwargs) + return track_tensor_tree(out, out_proxy, constant=None, tracer=proxy_mode.tracer) + + +@hints_wrapper.py_impl(ProxyTorchDispatchMode) +def inner(proxy_mode, body_fn, args, kwargs, hints): + if proxy_mode.enable_tracing: + return trace_hints_wrapper( + proxy_mode, hints_wrapper, body_fn, args, kwargs, hints + ) + else: + return hints_wrapper(body_fn, args, kwargs, hints) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/invoke_subgraph.py b/lib/python3.12/site-packages/torch/_higher_order_ops/invoke_subgraph.py new file mode 100644 index 0000000000000000000000000000000000000000..b192e551669e8ed94cd0f840335c4018949e3e63 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/invoke_subgraph.py @@ -0,0 +1,308 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs + + +from typing import List, Optional, Tuple, Union + +import torch +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._dispatch.python import suspend_functionalization +from torch._higher_order_ops.utils import ( + _from_fun, + _maybe_reenter_make_fx, + clone_outputs_aliasing_inputs, + get_dummy_aot_autograd_config, + prepare_fw_with_masks, + reenter_make_fx, + save_tensors_and_symints_for_backward, + saved_tensors_and_symints, +) +from torch._ops import HigherOrderOperator +from torch._subclasses import FakeTensorMode +from torch._subclasses.functional_tensor import disable_functional_mode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ProxyTorchDispatchMode, + track_tensor_tree, +) +from torch.fx.graph_module import GraphModule + + +invoke_subgraph_counter = 0 + + +class InvokeSubgraphHOP(HigherOrderOperator): + def __init__(self) -> None: + super().__init__("invoke_subgraph") + + # identifier is setup by upper part of the stack. This helps us in + # identifying two invoke_subgraph calls have same subgraph. + def __call__( + self, + subgraph: GraphModule, + identifier: Optional[str], + operands: Union[ + List[Union[torch.Tensor, int, torch.SymInt]], + Tuple[Union[torch.Tensor, int, torch.SymInt]], + ], + ): + assert identifier is None or isinstance( + identifier, str + ), "identifier must be a None or a string" + + assert isinstance( + operands, (list, tuple) + ), f"invoke_subgraph operands must be a list or tuple of tensors/ints/SymInts {operands}" + assert all( + isinstance(o, (torch.Tensor, int, torch.SymInt)) for o in operands + ), f"invoke_subgraph operands must be a list of tensors/ints/SymInts {operands}" + + return super().__call__(subgraph, identifier, operands) + + +invoke_subgraph = InvokeSubgraphHOP() + + +def invoke_subgraph_placeholder(subgraph, *args, **kwargs): + # Just a placeholder for Dynamo to replace with invoke_subgraph + return subgraph(*args, **kwargs) + + +def mark_compile_region(fn=None): + """ + This wrapper instructs torch.compile to compile the wrapped region once and + reuse the compiled artifact, instead of the usual way of aggressively + inlining the function. + + Under the hood, it tells TorchDynamo to use InvokeSubgraph HOP for the + region. For PyTorch eager, this is a no-op. + """ + + def wrap(func): + def inner(*args, **kwargs): + return invoke_subgraph_placeholder(func, *args, **kwargs) + + return inner + + if fn: + return wrap(fn) + else: + return wrap + + +def get_invoke_subgraph_cache(): + cache = None + if tracing_ctx := torch._guards.TracingContext.try_get(): + cache = tracing_ctx.hop_dispatch_set_cache.get_cache(invoke_subgraph) + return cache + + +def trace_joint_graph(fn, fw_inputs, fw_outputs): + """ + Naively trace out a joint graph. This simplifies the reconstruction of joint + graph in the min-cut partitioner later on. + """ + from torch._functorch.aot_autograd import create_joint + + dummy_aot_config = get_dummy_aot_autograd_config() + + # This joint_fn is inserted as the backward graph as is. This simplifies the + # min-cut partitioner work later on. + # Input signature - (*primals, *tangents) + # Output signature - (*grads, *fw_outs) + # The output signature is deliberately kept grads first and fw_outs second. + # Having grads first makes the min-cut partitioner HOP graph stitching + # easier. + def joint_fn(*primals_and_tangents): + primals = primals_and_tangents[: len(fw_inputs)] + tangents = primals_and_tangents[len(fw_inputs) :] + + fw_outs, grads = create_joint( + prepare_fw_with_masks(fn), aot_config=dummy_aot_config + )(primals, tangents) + + maybe_clone = clone_outputs_aliasing_inputs(primals_and_tangents) + + # return signature is deliberately kept (*grads, *fw_outs). This + # simplifies partitioning work later on. + return pytree.tree_map(maybe_clone, grads + list(fw_outs)) + + primals = list(fw_inputs) + # This assumes that the tangent strides match fw_outputs strides. Check the + # InvokeSubgraphAutogradOp backward op for the contiguous call. + tangents = [_from_fun(out) for out in fw_outputs] + + joint_operands = primals + tangents + + return _maybe_reenter_make_fx(joint_fn)(*joint_operands) + + +def create_fw_bw_graph(subgraph, operands, grad_outputs=None): + with suspend_functionalization(), disable_functional_mode(): + with disable_proxy_modes_tracing(): + # args are functional tensors, generate some example tensors + fw_inputs = pytree.tree_map(_from_fun, operands) + + if grad_outputs is None: + # Infer grad_outputs to be the same properties as the fw_outputs + # if they're not passed in. + grad_outputs = pytree.tree_map(_from_fun, subgraph(*fw_inputs)) + if any( + not isinstance(out, torch.Tensor) + for out in grad_outputs + if out is not None + ): + raise RuntimeError( + "Expect outputs of invoke_subgraph to only contains tensors or None. " + f"Got types {[type(out) for out in grad_outputs]}." + ) + + # Trace the forward subgraph + fw_graph = _maybe_reenter_make_fx(subgraph)(*fw_inputs) + + # Trace the joint graph and assign it to the bwd graph + bw_graph = trace_joint_graph( + subgraph, + fw_inputs, + grad_outputs, + ) + return fw_graph, bw_graph, len(grad_outputs) + + +class InvokeSubgraphAutogradOp(torch.autograd.Function): + """ + This autograd function op is to stash the backward graph in the ctx while + running forward. + """ + + @staticmethod + def forward(ctx, fw_graph, bw_graph, identifier, num_fw_outs, *operands): + ctx._fw_graph = fw_graph + ctx._bw_graph = bw_graph + ctx._identifier = identifier + ctx._num_fw_outs = num_fw_outs + + with torch._C._AutoDispatchBelowAutograd(): + out = invoke_subgraph( + fw_graph, + f"___forward_{identifier}", + operands, + ) + + save_tensors_and_symints_for_backward(ctx, operands) + return out + + @staticmethod + def backward(ctx, *grad_outs): + bw_graph = ctx._bw_graph + identifier = ctx._identifier + primals = saved_tensors_and_symints(ctx) + num_fw_outs = ctx._num_fw_outs + + # While tracing we made the assumption that tangents are contiguous. So, + # force the grad_outs to be contiguous. + contiguous_grad_outs = tuple([o.contiguous() for o in grad_outs]) + + # bw_graph is a joint graph with signature (*primals_and_tangents) and + # returns (*grads_and_fw_outs). To get the grads, we use the num_fw_outs + # to extract the grads. + primals_and_tangents = primals + contiguous_grad_outs + grads = invoke_subgraph( + bw_graph, f"___backward_{identifier}", primals_and_tangents + )[:-num_fw_outs] + return None, None, None, None, *grads + + +@invoke_subgraph.py_impl(DispatchKey.CompositeExplicitAutograd) +def _(subgraph, identifier, operands): + from torch.utils._python_dispatch import _get_current_dispatch_mode + + mode = _get_current_dispatch_mode() + assert mode is None, "Mode should never be enabled for CPU/CUDA key" + return subgraph(*operands) + + +@invoke_subgraph.py_impl(DispatchKey.Autograd) +def _(subgraph, identifier, operands): + if not torch.is_grad_enabled(): + with torch._C._AutoDispatchBelowAutograd(): + return invoke_subgraph(subgraph, identifier, operands) + + # A shortcut for the case where all inputs don't require gradient, + # we skip tracing the forward and backward graph. + if pytree.tree_all_only( + torch.Tensor, + lambda t: not t.requires_grad, # type: ignore[union-attr] + operands, + ): + with torch._C._AutoDispatchBelowAutograd(): + return invoke_subgraph(subgraph, identifier, operands) + + # Check if we have already traced the subgraph. + invoke_subgraph_cache = get_invoke_subgraph_cache() + if invoke_subgraph_cache: + if saved_autograd_fn := invoke_subgraph_cache.get_autograd_key_entry( + identifier + ): + return saved_autograd_fn(*operands) + + fw_graph, bw_graph, num_fw_outs = create_fw_bw_graph(subgraph, operands) + + def autograd_fn_callable(*args): + return InvokeSubgraphAutogradOp.apply( + fw_graph, bw_graph, identifier, num_fw_outs, *args + ) + + # Save the autograd_fn_callable in the dispatch set cache. + if invoke_subgraph_cache: + invoke_subgraph_cache.add_autograd_key_entry(identifier, autograd_fn_callable) + + return autograd_fn_callable(*operands) + + +@invoke_subgraph.py_functionalize_impl +def _(ctx, subgraph, identifier, operands): + unwrapped_operands = ctx.unwrap_tensors(operands) + with ctx.redispatch_to_next() as m: + # NB: There is an assumption that subgraph does not mutate inputs and + # there is no aliasing. Its Dynamo responsibility to prevent formation + # of invoke_subgraph ops if input aliasing/mutation is detected. + functionalized_subgraph = ctx.functionalize(subgraph) + out = invoke_subgraph(functionalized_subgraph, identifier, unwrapped_operands) + return ctx.wrap_tensors(out) + + +@invoke_subgraph.py_impl(FakeTensorMode) +def _(mode, subgraph, identifier, operands): + # TODO(anijain2305) - Implement fake tensor caching. + with mode: + return subgraph(*operands) + + +@invoke_subgraph.py_impl(ProxyTorchDispatchMode) +def _(proxy_mode: ProxyTorchDispatchMode, subgraph, identifier, operands): + # Check if we have already traced the subgraph. + graph = None + invoke_subgraph_cache = get_invoke_subgraph_cache() + if invoke_subgraph_cache: + graph = invoke_subgraph_cache.get_proxy_dispatch_entry(identifier) + + if graph is None: + graph = reenter_make_fx(subgraph)(*operands) + assert isinstance(proxy_mode.tracer, torch.fx.Tracer) + qualname = proxy_mode.tracer.get_fresh_qualname("repeated_subgraph") + proxy_mode.tracer.root.register_module(qualname, graph) + if invoke_subgraph_cache: + invoke_subgraph_cache.add_proxy_dispatch_entry(identifier, graph) + + node_args = (graph, identifier, operands) + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, node_args) # type: ignore[union-attr] + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", invoke_subgraph, proxy_args, {} + ) + + example_out = invoke_subgraph(graph, identifier, operands) + return track_tensor_tree( + example_out, out_proxy, constant=None, tracer=proxy_mode.tracer + ) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/map.py b/lib/python3.12/site-packages/torch/_higher_order_ops/map.py new file mode 100644 index 0000000000000000000000000000000000000000..cbd7f9c45fd4769bf94ecc33db2f80c567a293e7 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/map.py @@ -0,0 +1,264 @@ +# mypy: allow-untyped-defs +import torch +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._dispatch.python import suspend_functionalization +from torch._functorch.aot_autograd import AOTConfig, create_joint +from torch._higher_order_ops.utils import ( + _has_potential_branch_input_alias, + _has_potential_branch_input_mutation, + _maybe_run_with_interpreter, + reenter_make_fx, + UnsupportedAliasMutationException, +) +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch._subclasses.functional_tensor import disable_functional_mode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + make_fx, + ProxyTorchDispatchMode, + track_tensor_tree, +) + +from .utils import ( + _from_fun, + _stack_pytree, + _unstack_pytree, + clone_outputs_aliasing_inputs, + prepare_fw_with_masks, + save_tensors_and_symints_for_backward, + saved_tensors_and_symints, +) + + +# TODO: We add this to prevent dymamo from tracing into map_wrapper, +# remove the wrapper call when it's ready. +class MapWrapper(HigherOrderOperator): + def __init__(self): + super().__init__("map") + + def __call__(self, xs, *args): + return map_wrapper(xs, *args) + + +class MapImpl(HigherOrderOperator): + def __init__(self): + super().__init__("map_impl") + + def __call__(self, *args, **kwargs): + return super().__call__(*args, **kwargs) + + +map = MapWrapper() + +map_impl = MapImpl() + +dummy_aot_config = AOTConfig( + fw_compiler=None, # type: ignore[arg-type] + bw_compiler=None, # type: ignore[arg-type] + partition_fn=None, # type: ignore[arg-type] + decompositions={}, + num_params_buffers=0, + aot_id=0, + keep_inference_input_mutations=False, +) + + +def create_fw_bw_graph(f, num_mapped_args, *args): + mapped_xs = args[:num_mapped_args] + pos_args = args[num_mapped_args:] + + # See Note [HOP create fw_bw graph] in create_fw_bw_graph in utils.py + + with suspend_functionalization(), disable_functional_mode(): + with disable_proxy_modes_tracing(): + unwrapped_mapped_xs = pytree.tree_map(_from_fun, mapped_xs) + example_xs = _unstack_pytree(unwrapped_mapped_xs)[0] + + example_pos_args = [ + _from_fun(arg) if isinstance(arg, torch.Tensor) else arg + for arg in pos_args + ] + example_flat_out = pytree.tree_map( + _from_fun, f(*example_xs, *example_pos_args) + ) + if any( + not isinstance(out, torch.Tensor) + for out in example_flat_out + if out is not None + ): + raise RuntimeError( + "Expect outputs of map only contains tensors or None. " + f"Got types {[type(out) for out in example_flat_out]}." + ) + example_grad = [_from_fun(out) for out in example_flat_out] + + fw_graph = make_fx(f)(*example_xs, *example_pos_args) + + def joint_f(*example_args): + joint_mapped_args = example_args[:joint_num_mapped] + args = example_args[joint_num_mapped:] + + mapped_input = joint_mapped_args[:num_mapped_args] + mapped_grads = joint_mapped_args[num_mapped_args:] + + joint = create_joint(prepare_fw_with_masks(f), aot_config=dummy_aot_config) + _, grads = joint( + list(mapped_input) + list(args), + [ + grad + for grad in mapped_grads + if grad is not None and grad.requires_grad + ], + ) + + # In order to keep map functional for backward graph, + # we clone outputs that are aliasing inputs + maybe_clone = clone_outputs_aliasing_inputs(example_args) + + return pytree.tree_map(maybe_clone, grads) + + joint_num_mapped = len(example_grad) + len(example_xs) + joint_graph = make_fx(joint_f)(*example_xs, *example_grad, *example_pos_args) + return fw_graph, joint_graph + + +def map_wrapper(f, xs, *args): + flat_xs, xs_spec = pytree.tree_flatten(xs) + if not all(isinstance(t, torch.Tensor) for t in flat_xs): + raise RuntimeError(f"Mapped xs can only consist of tensors. Got xs {flat_xs}.") + + num_mapped_args = len(flat_xs) + shapes = [xs.shape for xs in flat_xs] + leading_dim_size = shapes[0][0] + if leading_dim_size == 0: + raise RuntimeError("Leading dimensions of mapped xs cannot be 0.") + + if any(cur_shape[0] != leading_dim_size for cur_shape in shapes): + raise RuntimeError( + f"Leading dimensions of mapped xs must be consistent. Got shapes {shapes}." + ) + + out_spec = None + + def flat_fn(*flat_args): + xs = pytree.tree_unflatten(list(flat_args[:num_mapped_args]), xs_spec) + unflattened_out = f(xs, *flat_args[num_mapped_args:]) + flat_out, tmp_out_spec = pytree.tree_flatten(unflattened_out) + + nonlocal out_spec + out_spec = tmp_out_spec + return flat_out + + return pytree.tree_unflatten( + map_impl(flat_fn, flat_xs, args), out_spec # type: ignore[arg-type] + ) + + +class MapAutogradOp(torch.autograd.Function): + @staticmethod + def forward(ctx, fw_graph, joint_graph, num_mapped_args, *flat_args): + save_tensors_and_symints_for_backward(ctx, flat_args) + ctx._joint_graph = joint_graph + ctx._num_mapped_args = num_mapped_args + with torch._C._AutoDispatchBelowAutograd(): + return ( + *map_impl( + fw_graph, flat_args[:num_mapped_args], flat_args[num_mapped_args:] + ), + ) + + @staticmethod + def backward(ctx, *flat_grads): + fw_args = saved_tensors_and_symints(ctx) + fw_mapped_args = fw_args[: ctx._num_mapped_args] + pos_args = fw_args[ctx._num_mapped_args :] + + grads = map_impl( + ctx._joint_graph, + fw_mapped_args + flat_grads, + pos_args, + ) + return None, None, None, *grads + + +def trace_map(proxy_mode, func_overload, f, xs, pos_args): + leading_dim_size = xs[0].shape[0] + + example_input = _unstack_pytree(xs)[0] + body_graph = f + + body_graph = reenter_make_fx(body_graph)(*example_input, *pos_args) + + next_name = proxy_mode.tracer.get_fresh_qualname("body_graph_") + + proxy_mode.tracer.root.register_module(next_name, body_graph) + + with disable_proxy_modes_tracing(): + example_outs = body_graph(*example_input, *pos_args) + + def expand_tensor(t): + if isinstance(t, torch.Tensor): + return t.expand(leading_dim_size, *t.shape) + return t + + expanded_outs = pytree.tree_map(expand_tensor, example_outs) + + node_args = (body_graph, list(xs), list(pos_args)) + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, node_args) + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", func_overload, proxy_args, {}, name="map_impl" + ) + return track_tensor_tree( + expanded_outs, out_proxy, constant=None, tracer=proxy_mode.tracer + ) + + +@map_impl.py_impl(DispatchKey.CompositeExplicitAutograd) +def map_dense(f, xs, pos_args): + pytrees = [f(*inp, *pos_args) for inp in _unstack_pytree(xs)] + return _stack_pytree(pytrees) + + +@map_impl.py_impl(DispatchKey.Autograd) +def map_autograd(f, xs, pos_args): + num_mapped_args = len(xs) + fw_graph, bw_graph = create_fw_bw_graph(f, num_mapped_args, *xs, *pos_args) + flat_out = MapAutogradOp.apply(fw_graph, bw_graph, num_mapped_args, *xs, *pos_args) + return flat_out + + +@map_impl.py_impl(ProxyTorchDispatchMode) +def map_proxy_torch_dispatch_mode(mode, f, xs, args): + return trace_map(mode, map_impl, f, xs, args) + + +@map_impl.py_impl(FakeTensorMode) +def map_fake_tensor_mode(mode, f, xs, args): + with mode: + return map_dense(f, xs, args) + + +@map_impl.py_functionalize_impl +def map_functionalize(ctx, f, xs, pos_args): + unwrapped_xs = ctx.unwrap_tensors(xs) + unwrapped_args = ctx.unwrap_tensors(pos_args) + wrapped_fn = ctx.functionalize(_maybe_run_with_interpreter(f)) + + with ctx.redispatch_to_next(): + with disable_proxy_modes_tracing(): + example_inputs = (*_unstack_pytree(unwrapped_xs)[0], *unwrapped_args) + pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch + if _has_potential_branch_input_mutation( + f, example_inputs, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException("torch.map is mutating the input!") + + if _has_potential_branch_input_alias( + f, example_inputs, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException("torch.map is aliasing the input!") + + map_return = map_impl(wrapped_fn, unwrapped_xs, unwrapped_args) + return ctx.wrap_tensors(map_return) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/out_dtype.py b/lib/python3.12/site-packages/torch/_higher_order_ops/out_dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..3a2dfbe8ae2fb5f394d507fce632fe9802261136 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/out_dtype.py @@ -0,0 +1,166 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs + +import torch +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._higher_order_ops.utils import autograd_not_implemented +from torch._ops import HigherOrderOperator +from torch._prims_common import elementwise_dtypes, ELEMENTWISE_TYPE_PROMOTION_KIND +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + maybe_handle_decomp, + ProxyTorchDispatchMode, + track_tensor_tree, +) + + +# TODO to figure out a more generic approach +ALLOWABLE_OPS = [ + torch.ops.aten.linear.default, + torch.ops.aten.mm.default, + torch.ops.aten.conv2d.default, + torch.ops.aten.convolution.default, + torch.ops.aten.mul.Tensor, + torch.ops.aten.mul.Scalar, + torch.ops.aten.div.Tensor, + torch.ops.aten.div.Scalar, +] + + +class OutDtypeOperator(HigherOrderOperator): + """ + The out_dtype operator takes an existing ATen functional operator, an + `out_dtype` argument, and arguments to the original operator, and executes + the original operator and returns a Tensor with the `out_dtype` precision. + This operator does not mandate a compute precision so it allows the + representation to not be opinionated about the exact implementation. + + The general implementation for all operators will be the following: + 1. Promote inputs dtypes based on default PyTorch dtype promotion rules, + using the dtypes of all input Tensors/Scalars and the `out_dtype` + arugument. + 2. Execute the operator + 3. Cast the output to `out_dtype` + """ + + def __init__(self) -> None: + super().__init__("out_dtype") + + def __call__(self, op, output_dtype, *args): + if not isinstance(op, torch._ops.OpOverload): + raise ValueError("out_dtype's first argument must be an OpOverload") + if op._schema.is_mutable: + raise ValueError( + "out_dtype's first argument needs to be a functional operator" + ) + if not ( + len(op._schema.returns) == 1 + and isinstance(op._schema.returns[0].type, torch.TensorType) + ): + raise ValueError( + "out_dtype's can only apply to ops that return a single tensor" + f"Instead got {[r.type for r in op._schema.returns]}" + ) + + if op not in ALLOWABLE_OPS: + raise ValueError( + f"out_dtype only allows the following operators: {ALLOWABLE_OPS}." + ) + + res = super().__call__(op, output_dtype, *args) + + return res + + +out_dtype = OutDtypeOperator() + + +def trace_out_dtype(proxy_mode, func_overload, op, output_dtype, *args): + # NB: Long-term we should put the decomposition logic into + # ProxyTorchDispatchMode so that people do not need to call maybe_handle_decomp + # in all HigherOrderOp proxy implementations. + r = maybe_handle_decomp(proxy_mode, func_overload, (op, output_dtype, *args), {}) + if r is not NotImplemented: + return r + + with disable_proxy_modes_tracing(): + # This is a simplified implementation of this operator just for tracing. + # Actual implementation may also first promote the arguments + out = op(*args).to(dtype=output_dtype) + + node_args = (op, output_dtype, *args) + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, node_args) + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", func_overload, proxy_args, {}, name="out_dtype" + ) + return track_tensor_tree(out, out_proxy, constant=None, tracer=proxy_mode.tracer) + + +@out_dtype.py_impl(DispatchKey.CompositeExplicitAutograd) +def out_dtype_dense(op: torch._ops.OpOverload, output_dtype: torch.dtype, *args): + if is_int_mm(op, output_dtype, args): + return torch._int_mm(*args) + return out_dtype_fallback(op, output_dtype, *args) + + +def is_int_mm(op, output_dtype, args): + return ( + op == torch.ops.aten.mm.default + and output_dtype == torch.int32 + and len(args) == 2 + and args[0].dtype == torch.int8 + and args[1].dtype == torch.int8 + and args[0].is_cuda + and args[1].is_cuda + ) + + +def out_dtype_fallback(op, output_dtype, *args): + flat_inputs = pytree.arg_tree_leaves(*args) + [torch.ones(1, dtype=output_dtype)] + promote_dtype: torch.dtype = elementwise_dtypes( + *flat_inputs, + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + )[0] + + casted_args = pytree.tree_map_only( + torch.Tensor, lambda arg: arg.to(dtype=promote_dtype), args + ) + res = op(*casted_args).to(dtype=output_dtype) + return res + + +out_dtype.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(out_dtype, deferred_error=True) +) + + +@out_dtype.py_impl(ProxyTorchDispatchMode) +def out_dtype_proxy( + mode: ProxyTorchDispatchMode, + op: torch._ops.OpOverload, + output_dtype: torch.dtype, + *args, +): + return trace_out_dtype(mode, out_dtype, op, output_dtype, *args) + + +@out_dtype.py_impl(FakeTensorMode) +def out_dtype_fake_tensor_mode( + mode: FakeTensorMode, + op: torch._ops.OpOverload, + output_dtype: torch.dtype, + *args, +): + with mode: + return out_dtype_dense(op, output_dtype, *args) + + +@out_dtype.py_functionalize_impl +def out_dtype_func(ctx, op, output_dtype, *args): + unwrapped_args = tuple(ctx.unwrap_tensors(arg) for arg in args) + + with ctx.redispatch_to_next(): + res = out_dtype(op, output_dtype, *unwrapped_args) + return ctx.wrap_tensors(res) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/prim_hop_base.py b/lib/python3.12/site-packages/torch/_higher_order_ops/prim_hop_base.py new file mode 100644 index 0000000000000000000000000000000000000000..04c52496b385f318dc9c2ce4fedf04b098646073 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/prim_hop_base.py @@ -0,0 +1,190 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs + +import abc + +import torch +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._dispatch.python import suspend_functionalization +from torch._higher_order_ops.utils import reenter_make_fx +from torch._ops import HigherOrderOperator +from torch._subclasses import FakeTensorMode +from torch._subclasses.functional_tensor import disable_functional_mode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ProxyTorchDispatchMode, + track_tensor_tree, +) + + +class PrimHOPBase(HigherOrderOperator, abc.ABC): + """ + This is the "Base" HOP implementation for a HOP that looks like: + + call_subgraph_hop(subgraph, operands, **kwargs) + + That is: + 1) the HOP is a "prim" (it stays alive until Inductor) + 2) the HOP's semantics are subgraph(*operands) + + To use this, please subclass this class and override methods as necessary: + ``` + class InvokeQuant(PrimHOPBase): + def __init__(self): + return super().__init__("invoke_quant") + + invoke_quant = InvokeQuant() + + def g(x): + return x.sin().cos() + + @torch.compile(backend="aot_eager") + def f(x): + return invoke_quant(g, (x,), scheme="nf4") + ``` + + NOTE: don't subclass PrimHOPBase out of tree! That is not allowed. All + usages must be in tree. + """ + + def __init__(self, hop_name) -> None: + super().__init__(hop_name) + + # Set up the registrations + # If you want to override any of these, override them in your subclass. + self.py_impl(DispatchKey.Autograd)(self._call_Autograd) + self.py_functionalize_impl(self._call_Functionalize) + self.py_impl(ProxyTorchDispatchMode)(self._call_ProxyTorchDispatchMode) + self.py_impl(FakeTensorMode)(self._call_FakeTensorMode) + self.py_impl(DispatchKey.CompositeExplicitAutograd)( + self._call_CompositeExplicitAutograd + ) + + def __call__(self, subgraph, operands, *unused, **kwargs): + # We accept *unused (and *_) to make mypy happy. Otherwise mypy + # complains that we're violating LSP. We are violating LSP, but it's + # OK for the purposes of implementation-sharing (end users should never + # subclass these methods; only in-tree PyTorch developers are allowed to). + assert len(unused) == 0 + if not isinstance(subgraph, (torch.fx.GraphModule, FunctionWithNoFreeVars)): + raise RuntimeError( + f"{self._name}: when calling this API without torch.compile, " + f"we require that the subgraph be a torch.fx.GraphModule (or " + f"a function we know doesn't have free variables)." + ) + return super().__call__(subgraph, operands, **kwargs) + + def _call_Autograd(self, subgraph, operands, *_, **kwargs): + if isinstance(subgraph, torch.fx.GraphModule): + pass + if not torch.is_grad_enabled() or pytree.tree_all_only( + torch.Tensor, + lambda t: not t.requires_grad, # type: ignore[union-attr] + operands, + ): + with torch._C._AutoDispatchBelowAutograd(): + return self(subgraph, operands, **kwargs) + + # We assume the subgraph doesn't mutate inputs and there is no aliasing. + # In the PT2 stack, this is Dynamo's responsibility to figure out. + return PrimHOPBaseFunction.apply(self, subgraph, kwargs, *operands) + + def _call_CompositeExplicitAutograd(self, subgraph, operands, *_, **kwargs): + from torch.utils._python_dispatch import _get_current_dispatch_mode + + mode = _get_current_dispatch_mode() + assert mode is None, "Mode should never be enabled for CPU/CUDA key" + return subgraph(*operands) + + def _call_ProxyTorchDispatchMode( + self, proxy_mode, subgraph, operands, *_, **kwargs + ): + traced_graph = reenter_make_fx(subgraph)(*operands) + assert isinstance(proxy_mode.tracer, torch.fx.Tracer) + qualname = proxy_mode.tracer.get_fresh_qualname("subgraph") + proxy_mode.tracer.root.register_module(qualname, traced_graph) + + node_args = (traced_graph, operands) + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, node_args) # type: ignore[attr-defined] + proxy_kwargs = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, kwargs) # type: ignore[attr-defined] + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", self, proxy_args, proxy_kwargs + ) + + out = self(subgraph, operands, **kwargs) + return track_tensor_tree( + out, out_proxy, constant=None, tracer=proxy_mode.tracer # type: ignore[arg-type] + ) + + def _call_FakeTensorMode(self, mode, subgraph, operands, *_, **kwargs): + # TODO: this should probably route through FakeTensorMode to reuse caching + with mode: + return subgraph(*operands) + + def _call_Functionalize(self, ctx, subgraph, operands, *_, **kwargs): + unwrapped_operands = ctx.unwrap_tensors(operands) + with ctx.redispatch_to_next() as m: + # We assume the subgraph doesn't mutate inputs and there is no aliasing. + # In the PT2 stack, this is Dynamo's responsibility to figure out. + functionalized_subgraph = FunctionWithNoFreeVars( + ctx.functionalize(subgraph) + ) + out = self(functionalized_subgraph, unwrapped_operands, **kwargs) + return ctx.wrap_tensors(out) + + +class PrimHOPBaseFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, hop, subgraph, kwargs, *operands): + ctx.hop = hop + ctx.operands = operands + ctx.subgraph = subgraph + ctx.kwargs = kwargs + + with torch._C._AutoDispatchBelowAutograd(): + return hop(subgraph, operands, **kwargs) + + @staticmethod + def backward(ctx, *grad_outputs): + subgraph = ctx.subgraph + operands = ctx.operands + kwargs = ctx.kwargs + + # TODO: Something special needs to happen with min cut partitioner + with suspend_functionalization(), disable_functional_mode(), torch.enable_grad(): + with disable_proxy_modes_tracing(): + from .invoke_subgraph import create_fw_bw_graph + from .utils import _from_fun + + fw_inputs = pytree.tree_map(_from_fun, operands) + fw_outputs = subgraph(*fw_inputs) + _, joint_graph, _ = create_fw_bw_graph( + subgraph, fw_inputs, grad_outputs + ) + + # The joint graph returns (*grad_inputs, *fwd_outputs). + # We only need the grad_inputs. + def bwd_fn(*args): + operands = args[: -len(grad_outputs)] + grad_outs = args[-len(grad_outputs) :] + result = joint_graph(*operands, *grad_outs) + grad_inputs = result[: -len(grad_outputs)] + return grad_inputs + + return ( + None, + None, + None, + *ctx.hop( + FunctionWithNoFreeVars(bwd_fn), (*operands, *grad_outputs), **kwargs + ), + ) + + +class FunctionWithNoFreeVars: + def __init__(self, fn): + self.fn = fn + + def __call__(self, *args, **kwargs): + return self.fn(*args, **kwargs) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/run_const_graph.py b/lib/python3.12/site-packages/torch/_higher_order_ops/run_const_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..1f49ee28394a1c464612d2c5ee8a59abd6be1b49 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/run_const_graph.py @@ -0,0 +1,60 @@ +# mypy: allow-untyped-defs +import torch +from torch._C import DispatchKey +from torch._higher_order_ops.utils import autograd_not_implemented +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode, track_tensor_tree +from torch.utils import _pytree as pytree + + +class RunConstGraph(HigherOrderOperator): + def __init__(self): + super().__init__("run_const_graph") + + def __call__(self, *args): + return super().__call__(*args) + + +run_const_graph = RunConstGraph() + + +@run_const_graph.py_impl(ProxyTorchDispatchMode) +def run_const_graph_dispatch_mode(mode, *args): + const_gm, weights = args + p_args = pytree.tree_map(mode.tracer.unwrap_proxy, args) + assert isinstance(const_gm, torch.fx.GraphModule) + assert not hasattr(mode.tracer.root, "_const_graph") + mode.tracer.root.register_module("_const_graph", const_gm) + + proxy = mode.tracer.create_proxy("call_function", run_const_graph, p_args, {}) + + out = const_gm(*weights) + return track_tensor_tree(out, proxy, constant=None, tracer=mode.tracer) + + +@run_const_graph.py_functionalize_impl +def run_const_graph_functional(ctx, *args): + unwrapped_args = ctx.unwrap_tensors(args) + + with ctx.redispatch_to_next(): + out = run_const_graph(*unwrapped_args) + return ctx.wrap_tensors(out) + + +run_const_graph.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(run_const_graph, deferred_error=True) +) + + +@run_const_graph.py_impl(FakeTensorMode) +def run_const_graph_fake_tensor_mode(mode, graph, args): + assert isinstance(graph, torch.fx.GraphModule) + with mode: + return graph(*args) + + +@run_const_graph.py_impl(DispatchKey.CPU) +def run_const_graph_cpu(graph, args): + assert isinstance(graph, torch.fx.GraphModule) + return graph(*args) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/scan.py b/lib/python3.12/site-packages/torch/_higher_order_ops/scan.py new file mode 100644 index 0000000000000000000000000000000000000000..c151941178d9db3a0d6484ef9751f4d76cd501cf --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/scan.py @@ -0,0 +1,518 @@ +# mypy: allow-untyped-defs +import functools +import itertools +from typing import Any, Callable, List, Tuple + +import torch +import torch._prims_common as utils +import torch._subclasses.functional_tensor +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._higher_order_ops.utils import ( + _has_potential_branch_input_alias, + _has_potential_branch_input_mutation, + _set_compilation_env, + autograd_not_implemented, + reenter_make_fx, + unique_graph_id, + UnsupportedAliasMutationException, + validate_subgraph_args_types, +) +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + _temp_remove_metadata_torch_function_mode, + disable_proxy_modes_tracing, + ProxyTorchDispatchMode, + track_tensor_tree, +) +from torch.utils._python_dispatch import _get_current_dispatch_mode + + +aten = torch._ops.ops.aten + + +def wrap_combine_fn_flat( + *args, combine_fn, spec_init, spec_xs, num_init_leaves, num_inp_leaves +): + assert len(args) == (num_init_leaves + num_inp_leaves) + carry = pytree.tree_unflatten(args[:num_init_leaves], spec_init) + xs = pytree.tree_unflatten(args[num_init_leaves:], spec_xs) + carry, combined = combine_fn(carry, xs) + carry_flat = pytree.tree_leaves(carry) + combined_flat = pytree.tree_leaves(combined) + assert num_init_leaves == len(carry_flat) + return [*carry_flat, *combined_flat] + + +def _extract_carry_and_out(flat_out: List[Any], num_carry: int): + return flat_out[:num_carry], flat_out[num_carry:] + + +def scan( + combine_fn: Callable[ + [pytree.PyTree, pytree.PyTree], Tuple[pytree.PyTree, pytree.PyTree] + ], + init: pytree.PyTree, + xs: pytree.PyTree, + *, + dim: int = 0, + reverse: bool = False, +) -> Tuple[pytree.PyTree, pytree.PyTree]: + r""" + Performs an inclusive scan with a combine function. + + .. warning:: + `torch.scan` is a prototype feature in PyTorch. It currently + does not support autograd and you may run into miscompiles. + Read more about feature classification at: + https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype + + Args: + combine_fn (Callable): A binary callable with type ``(Tensor, Tensor) -> (Tensor, Tensor)``, + or if xs is a pytree ``(pytree, pytree) -> (pytree, pytree)``. + The first input to ``combine_fn`` is the previous or initial scan carry + and the second input element to ``combine_fn`` is a slice of the input along dim. + The first output element of ``combine_fn`` is the next scan carry + and the second output of ``combine_fn`` represents a slice of the output. + This function must be pure, i.e., no lifted arguments are supported at the moment + and may not have any side effects. + init (torch.Tensor or pytree with tensor leaves): The inital scan carry, a tensor, or nested pytree of tensors. + The ``init`` is expected to have the same pytree structure as the first output element (i.e. carry) + of ``combine_fn``. + xs (torch.Tensor or pytree with tensor leaves): The input tensor, or nested pytree of tensors. + + Kwargs: + dim (int): the dimension to scan over, default 0. + reverse (bool): A boolean stating if the scan should be reversed with respect to ``dim``, default ``False``. + + Returns: + final_carry (torch.Tensor or pytree with tensor leaves), + the final carry of the scan operation with same pytree structure as init. + out (torch.Tensor or pytree with tensor leaves), + each tensor leaf is a stacked output along first dim, where each slice is the output of a scan iteration. + + Example:: + + def add(x: torch.Tensor, y: torch.Tensor): + next_carry = y = x + y + return next_carry, y + + i0 = torch.zeros(1) + xs = torch.arange(5) + # returns torch.tensor([10.]), torch.tensor([[0], [1.], [3.], [6.], [10.]]) + last_carry, cumsum = scan(add, init=i0, xs=xs) + + + """ + if not callable(combine_fn): + raise RuntimeError("Combine_fn must be a callable, but got {combine_fn}") + if not isinstance(dim, int): + raise RuntimeError("Dim must be an int, but got " + str(type(dim))) + if not isinstance(reverse, bool): + raise RuntimeError("Reverse must be a bool, but got " + str(type(reverse))) + + leaves_init, spec_init = pytree.tree_flatten(init) + leaves_xs, spec_xs = pytree.tree_flatten(xs) + + if len(leaves_init) == 0: + raise RuntimeError("Init tensors must be provided") + for x in leaves_init: + if not isinstance(x, torch.Tensor): + raise RuntimeError(f"All init leaves must be a Tensor but got {x}") + for x in leaves_xs: + if not isinstance(x, torch.Tensor): + raise RuntimeError(f"All xs leaves must be a Tensor but got {x}") + if x.shape[dim] == 0: + raise RuntimeError( + f"All xs leaves must have a scan dimension > 0 but got {x}" + ) + + if len(leaves_xs) == 0: + return pytree.tree_unflatten(leaves_init, spec_init), xs + + shape = leaves_xs[0].shape + ndim = len(shape) + dim = utils.canonicalize_dim(ndim, dim) + + out = combine_fn( + pytree.tree_unflatten(leaves_init, spec_init), + pytree.tree_unflatten([elem.select(dim, 0) for elem in leaves_xs], spec_xs), + ) + + # The first output needs to have the same pytree as init + carry_leaves = pytree.tree_leaves(out[0]) + if len(carry_leaves) != len(leaves_init): + raise RuntimeError( + f"The number of leaves of the pytree of the new carry produced by the operator is {len(carry_leaves)}\ +doesn't match the length of the pytree of the init {len(leaves_init)}" + ) + + def _check_new_carry_match_init(leaves_init, carry_leaves): + for i, (init, new_carry) in enumerate(zip(leaves_init, carry_leaves)): + if init.shape != new_carry.shape: + raise RuntimeError( + f"The shape of the new_carry[{i}] {new_carry.shape} doesn't match that of the init[{i}] {init.shape}." + ) + if init.stride() != new_carry.stride(): + raise RuntimeError( + f"The stride of the new_carry[{i}] {new_carry.stride()} doesn't match that of the init[{i}] {init.stride()}." + ) + if init.dtype != new_carry.dtype: + raise RuntimeError( + f"The dtype of the new_carry[{i}] {new_carry.dtype} doesn't match that of the init[{i}] {init.dtype}." + ) + if init.requires_grad != new_carry.requires_grad: + raise RuntimeError( + f"The requires_grad of the new_carry[{i}] {new_carry.requires_grad} doesn't match that of the init[{i}] {init.requires_grad}." # noqa: B950 + ) + + _check_new_carry_match_init(leaves_init, carry_leaves) + + # There are no pytree restrictions on the second output of the operator + out_leaves, tree_out = pytree.tree_flatten(out[1]) + + # TODO: Support closures/nn_modules in order to be able represent RNNs with scan + # TODO: Support _inductor lowering + # TODO: Support Autograd + # TODO: Unify handling of pytrees for control flow ops, such as cond, while_loop, etc. + # TODO: Unify the list inputs of control flow ops to tuple. + + combine_fn = functools.partial( + wrap_combine_fn_flat, + combine_fn=combine_fn, + spec_init=spec_init, + spec_xs=spec_xs, + num_init_leaves=len(leaves_init), + num_inp_leaves=len(leaves_xs), + ) + + def run_flattened_scan(combine_fn, leaves_init, leaves_xs, dim, reverse): + return scan_op( + combine_fn, leaves_init, leaves_xs, dim, reverse, additional_inputs=[] + ) + + if not torch._dynamo.is_compiling(): + from torch._dynamo.backends.debugging import ( + make_eager_backend_with_torch_function_mode, + ) + + with _set_compilation_env(), torch._dynamo.utils.disable_cache_limit(): + with _temp_remove_metadata_torch_function_mode() as metadata_mode: + if metadata_mode: + backend = make_eager_backend_with_torch_function_mode(metadata_mode) + else: + backend = "eager" + result = torch.compile( + run_flattened_scan, backend=backend, fullgraph=True + )( + combine_fn, + leaves_init, + leaves_xs, + dim=dim, + reverse=reverse, + ) + else: + result = run_flattened_scan(combine_fn, leaves_init, leaves_xs, dim, reverse) + + result_carry, result_flat = _extract_carry_and_out( + result, + len(leaves_init), + ) + + return pytree.tree_unflatten(result_carry, spec_init), pytree.tree_unflatten( + result_flat, tree_out + ) + + +class ScanOp(HigherOrderOperator): + def __init__(self): + super().__init__("scan") + + def __call__(self, combine_fn, init, xs, dim, reverse, additional_inputs): + assert isinstance(additional_inputs, list), "additional_inputs must be a list." + validate_subgraph_args_types(additional_inputs) + return super().__call__(combine_fn, init, xs, dim, reverse, additional_inputs) + + +scan_op = ScanOp() + + +def generic_scan(operator, init, xs, dim=0, reverse=False, additional_inputs=None): + additional_inputs = additional_inputs if additional_inputs is not None else [] + + def _scan(init, xs): + """Perform scan on `elems` using `elems_init.""" + carry = init + if len(xs) == 0: + return carry, [] + + num_elems = xs[0].shape[dim] + if reverse: + ind = num_elems - 1 + else: + ind = 0 + + # Compute dummy shapes for the pre-allocation + num_init_leaves = len(init) + dummy_carry, dummy_out = _extract_carry_and_out( + operator( + *carry, + *[first_slice_copy(elem, dim) for elem in xs], + *additional_inputs, + ), + num_init_leaves, + ) + + # Pre-alocate + # outs -> Output matrix + # idxs -> Index matrix for scatter_ + # out: (num_elems, M, N, ...) + # idx: (1, M, N) + outs, idxs = zip( + *[ + [ + torch.zeros( + [num_elems] + list(e.size()), + dtype=e.dtype, + device=e.device, + ), + torch.ones_like(e, dtype=torch.int64).unsqueeze(0), + ] + for i, e in enumerate(dummy_out) + ] + ) + + def store_out_in_outs(out, ind): + # Store the intermediate out in the outs matrix + for o, x, idx in zip(outs, out, idxs): + # o: (num_elems, M, N ...) + # x: (M, N, ...) -> (1, M, N) + # ind * idx: (1, M, N,) with values to be ind + # essentially: o[ind][n][k] = x[0][n][k] + o.scatter_(0, ind * idx, x.unsqueeze(0)) + + for i in range(num_elems): + ind = i if not reverse else num_elems - i - 1 + carry, out = _extract_carry_and_out( + operator( + *carry, + *[elem.select(dim, ind) for elem in xs], + *additional_inputs, + ), + num_init_leaves, + ) + + # Store the inits in the outs matrix. + store_out_in_outs(out, ind) + + return [*carry, *list(outs)] + + scans = _scan(init, xs) + return scans + + +def first_slice_copy(t: torch.Tensor, dim: int) -> torch.Tensor: + return torch.select_copy(t, dim, 0) + + +# We also do a clone with contiguous_format. This is to be consistent with +# eager semantic of scan, which stacks the outputs. The result is contiguous +# as a result of the stack operation. +def stack_y(y: torch.Tensor, scan_length: int) -> torch.Tensor: + return ( + y.unsqueeze(0) + .repeat(*([scan_length] + [1] * y.ndim)) + .clone(memory_format=torch.contiguous_format) + ) + + +def trace_scan( + proxy_mode, + func_overload, + combine_fn: Callable, + init: List[torch.Tensor], + xs: List[torch.Tensor], + dim: int, + reverse: bool, + additional_inputs: List[torch.Tensor], +): + from torch._dynamo.utils import clone_input + + with disable_proxy_modes_tracing(): + sample_inits = [clone_input(x_init) for x_init in init] + sample_inputs = [first_slice_copy(x, dim) for x in xs] + sample_additional_inputs = [ + clone_input(x) if isinstance(x, torch.Tensor) else x + for x in additional_inputs + ] + combine_graph = reenter_make_fx(combine_fn)( + *sample_inits, *sample_inputs, *sample_additional_inputs + ) + + outputs = None + for node in combine_graph.graph.nodes: + if node.op == "output": + assert outputs is None + assert len(node.args) == 1 + outputs = node.args[0] + + assert outputs is not None + + carry, output = _extract_carry_and_out(outputs, len(init)) + + for ini, ca in zip(init, carry): + ini_meta = ini + carry_meta = ca.meta["tensor_meta"] + carry_val = ca.meta["val"] + if ( + carry_val.device != ini_meta.device + or carry_meta.dtype != ini_meta.dtype + or carry_meta.shape != ini_meta.shape + ): + raise RuntimeError( + f"Expected metadata of the combine_fn result {carry_meta} to be the same as " + + f"the metadata of init with {ini_meta}" + ) + + _, combine_graph_name = unique_graph_id(proxy_mode, prefix="scan_combine_graph") + + proxy_mode.tracer.root.register_module(combine_graph_name, combine_graph) + + args = (combine_graph, init, xs, dim, reverse, additional_inputs) + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, args) + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", func_overload, proxy_args, {}, name="scan" + ) + + with disable_proxy_modes_tracing(): + scan_length = xs[0].shape[dim] + fake_carry, fake_outputs = _extract_carry_and_out( + [o.meta["val"] for o in outputs], len(init) + ) + out = ( + *fake_carry, + *(stack_y(t, scan_length) for t in fake_outputs), + ) + + return track_tensor_tree(out, out_proxy, constant=None, tracer=proxy_mode.tracer) + + +@scan_op.py_impl(DispatchKey.CompositeExplicitAutograd) +def scan_op_dense(combine_fn, init, xs, dim, reverse, additional_inputs): + mode = _get_current_dispatch_mode() + assert mode is None, "Mode should never be enabled for CPU/CUDA key" + return generic_scan(combine_fn, init, xs, dim, reverse, additional_inputs) + + +scan_op.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(scan_op, deferred_error=True) +) + + +@scan_op.py_impl(ProxyTorchDispatchMode) +def scan_proxy_mode(mode, combine_fn, init, xs, dim, reverse, additional_inputs): + return trace_scan( + mode, scan_op, combine_fn, init, xs, dim, reverse, additional_inputs + ) + + +@scan_op.py_impl(FakeTensorMode) +def scan_fake_tensor_mode(mode, combine_fn, init, xs, dim, reverse, additional_inputs): + with mode: + scan_length = xs[0].shape[dim] + carry, outputs = _extract_carry_and_out( + combine_fn( + *init, + *[first_slice_copy(inp, dim) for inp in xs], + *additional_inputs, + ), + len(init), + ) + out = ( + *carry, + *(stack_y(t, scan_length) for t in outputs), + ) + return out + + +@scan_op.py_functionalize_impl +def scan_functionalize(ctx, combine_fn, init, xs, dim, reverse, additional_inputs): + unwrapped_xs = ctx.unwrap_tensors(xs) + unwrapped_init = ctx.unwrap_tensors(init) + unwrapped_additional_inputs = ctx.unwrap_tensors(additional_inputs) + with ctx.redispatch_to_next() as m: + functional_combine_fn = ctx.functionalize(combine_fn) + pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch + sample_unwrapped_xs_sliced = [ + first_slice_copy(inp, dim) for inp in unwrapped_xs + ] + sample_inputs = list( + itertools.chain( + unwrapped_init, + sample_unwrapped_xs_sliced, + unwrapped_additional_inputs, + ) + ) + if _has_potential_branch_input_mutation( + functional_combine_fn, sample_inputs, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException( + "Combine_fn might be modifying the input!" + ) + if _has_potential_branch_input_alias( + functional_combine_fn, sample_inputs, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException( + "Combine_fn might be aliasing the input!" + ) + ret = scan_op( + functional_combine_fn, + unwrapped_init, + unwrapped_xs, + dim, + reverse, + unwrapped_additional_inputs, + ) + return ctx.wrap_tensors(ret) + + +# dense implementation for scan. Used for testing only. +def _fake_scan(combine_fn, init, xs=None, dim=0, reverse=False): + carry_leaves, carry_spec = pytree.tree_flatten(init) + inp_leaves, inp_spec = pytree.tree_flatten(xs) + if xs is None or len(inp_leaves) == 0: + return init, [] + result_flat = [] + carry = carry_leaves + op = reversed if reverse else lambda x: x + + dummy_carry, dummy_out = combine_fn( + pytree.tree_unflatten(carry, carry_spec), + pytree.tree_unflatten( + [first_slice_copy(elem, dim) for elem in inp_leaves], + inp_spec, + ), + ) + dummy_out_leaves, dummy_out_spec = pytree.tree_flatten(dummy_out) + num_leaves = len(dummy_out_leaves) + + for ind in op(range(inp_leaves[0].size(dim))): + xs = [elem.select(dim, ind) for elem in inp_leaves] + + carry, y = combine_fn( + pytree.tree_unflatten(carry, carry_spec), + pytree.tree_unflatten(xs, inp_spec), + ) + carry, _ = pytree.tree_flatten(carry) + y, _ = pytree.tree_flatten(y) + result_flat.append(y) + + results = [ + torch.stack([e[leave_ind] for e in op(result_flat)]) + for leave_ind in range(num_leaves) + ] + return ( + pytree.tree_unflatten(carry, carry_spec), + pytree.tree_unflatten(results, dummy_out_spec), + ) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/strict_mode.py b/lib/python3.12/site-packages/torch/_higher_order_ops/strict_mode.py new file mode 100644 index 0000000000000000000000000000000000000000..e8543412c53a34e0f18653bdbbe07fe026daec3e --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/strict_mode.py @@ -0,0 +1,108 @@ +# mypy: allow-untyped-defs +import torch +import torch._subclasses.functional_tensor +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._functorch.utils import exposed_in +from torch._higher_order_ops.utils import _set_compilation_env, autograd_not_implemented +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + _temp_remove_metadata_torch_function_mode, + _temp_remove_pre_dispatch_torch_function_mode, + disable_proxy_modes_tracing, + make_fx, + ProxyTorchDispatchMode, + track_tensor_tree, +) +from torch.utils._python_dispatch import _get_current_dispatch_mode + + +@exposed_in("torch") +def strict_mode(callable, operands): + from torch._dynamo.backends.debugging import ( + make_eager_backend_with_torch_function_modes, + ) + + if torch.compiler.is_dynamo_compiling(): + return strict_mode_op(callable, operands) + + with _set_compilation_env(): + with _temp_remove_metadata_torch_function_mode() as metadata_mode: + with _temp_remove_pre_dispatch_torch_function_mode() as predispatch_mode: + modes = [metadata_mode, predispatch_mode] + modes = [mode for mode in modes if mode is not None] + if modes: + backend = make_eager_backend_with_torch_function_modes(modes) + else: + backend = "eager" + with torch._dynamo.utils.disable_cache_limit(): + return torch.compile( + strict_mode_op, backend=backend, fullgraph=True + )(callable, operands) + + +class StrictMode(HigherOrderOperator): + def __init__(self): + super().__init__("strict_mode") + + def __call__(self, callable, operands): + return super().__call__(callable, operands) + + +strict_mode_op = StrictMode() + + +@strict_mode_op.py_impl(DispatchKey.CompositeExplicitAutograd) +def strict_mode_op_dense(callable, operands): + mode = _get_current_dispatch_mode() + assert mode is None, "Mode should never be enabled for CPU/CUDA key" + return callable(*operands) + + +strict_mode_op.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(strict_mode_op, deferred_error=True) +) + + +@strict_mode_op.py_impl(ProxyTorchDispatchMode) +def inner(mode, callable, operands): + return trace_strict_mode(mode, strict_mode_op, callable, operands) + + +def trace_strict_mode(mode, strict_mode_op, callable, operands): + pre_dispatch = getattr(mode, "pre_dispatch", False) + + with disable_proxy_modes_tracing(): + graph = make_fx(callable, pre_dispatch=pre_dispatch)(*operands) + + graph_name = mode.tracer.get_fresh_qualname("strict_graph_") + mode.tracer.root.register_module(graph_name, graph) + + args = (graph, operands) + + proxy_args = pytree.tree_map(mode.tracer.unwrap_proxy, args) + + out_proxy = mode.tracer.create_proxy( + "call_function", strict_mode_op, proxy_args, {}, name="strict_mode" + ) + + out = graph(*operands) + return track_tensor_tree(out, out_proxy, constant=None, tracer=mode.tracer) + + +@strict_mode_op.py_impl(FakeTensorMode) +def strict_mode_fake_tensor_mode(mode, callable, operands): + with mode: + true_outs = callable(*operands) + return true_outs + + +@strict_mode_op.py_functionalize_impl +def strict_mode_func(ctx, callable, inputs): + unwrapped_inputs = ctx.unwrap_tensors(inputs) + with ctx.redispatch_to_next(): + functional_callable = ctx.functionalize(callable) + + cond_return = strict_mode_op(functional_callable, unwrapped_inputs) + return ctx.wrap_tensors(cond_return) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/torchbind.py b/lib/python3.12/site-packages/torch/_higher_order_ops/torchbind.py new file mode 100644 index 0000000000000000000000000000000000000000..b35b8d5b296d118045347c35940454d88da07133 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/torchbind.py @@ -0,0 +1,142 @@ +# mypy: allow-untyped-defs +import logging +from contextlib import contextmanager + +import torch +from torch._C import DispatchKey # @manual +from torch._functorch._aot_autograd.utils import KNOWN_TYPES +from torch._higher_order_ops.utils import autograd_not_implemented +from torch._library.fake_class_registry import _ns_and_class_name, FakeScriptObject +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode, track_tensor_tree +from torch.fx.node import has_side_effect +from torch.utils import _pytree as pytree + + +log = logging.getLogger(__name__) + + +# The call_torchbind operator represents a method invocation on a torchbind +# object. The calling convention is: +# call_torchbind(self: ScriptObject, method_name: str, *method_args, **method_kwargs) +# We do not expect users to write this operator directly. Instead it will be +# emitted by Dynamo when tracing encounters a torchbind object. +class CallTorchBind(HigherOrderOperator): + def __init__(self): + super().__init__("call_torchbind") + + def __call__(self, obj, method, *args, **kwargs): + return super().__call__(obj, method, *args, **kwargs) + + +call_torchbind = CallTorchBind() + +# Register this operator as side-effectful with FX. +# TODO: this is not really sufficient. While passes (hopefully) check +# Node.is_impure() and make good decisions, we also assume we can execute the +# graph as many times as we want without changing behavior, which is NOT true of +# ops that mutate torchbind object state. +has_side_effect(call_torchbind) + +_orig_scriptmethod_call = torch.ScriptMethod.__call__ + + +def torchbind_method_redispatch(self, *args, **kwargs): + if isinstance(self.raw_owner, torch.ScriptObject): + return call_torchbind(self.raw_owner, self.name, *args, **kwargs) + return _orig_scriptmethod_call(self, *args, **kwargs) + + +@contextmanager +def enable_torchbind_tracing(): + """Context manager that acts as a feature flag to enable torchbind tracing + behavior. Once torchbind tracing has been stabilized, we can remove this and + turn it always on. + """ + try: + KNOWN_TYPES.append(torch.ScriptObject) + torch.ScriptMethod.__call__ = torchbind_method_redispatch # type: ignore[method-assign] + yield + finally: + assert ( + KNOWN_TYPES.pop() is torch.ScriptObject + ), "Someone else messed with KNOWN_TYPES during tracing, exploding." + torch.ScriptMethod.__call__ = _orig_scriptmethod_call # type: ignore[method-assign] + + +@call_torchbind.py_impl(DispatchKey.CompositeExplicitAutograd) +def call_torchbind_impl(obj, method, *args, **kwargs): + if isinstance(obj, torch.ScriptObject): + return _orig_scriptmethod_call(getattr(obj, method), *args, **kwargs) + elif isinstance(obj, FakeScriptObject): + return getattr(obj.wrapped_obj, method)(*args, **kwargs) + else: + raise RuntimeError(f"Unsupported first arg type {type(obj)} for call_torchbind") + + +@call_torchbind.py_impl(ProxyTorchDispatchMode) +def inner(mode, *args, **kwargs): + proxy_args = pytree.tree_map(mode.tracer.unwrap_proxy, args) + proxy_kwargs = pytree.tree_map(mode.tracer.unwrap_proxy, kwargs) + + out_proxy = mode.tracer.create_proxy( + "call_function", + call_torchbind, + proxy_args, + proxy_kwargs, + ) + out = call_torchbind(*args, **kwargs) + + obj, method, *rest_args = args + if isinstance(obj, torch.ScriptObject): + ns, class_name = _ns_and_class_name( + obj._type().qualified_name() # type: ignore[attr-defined] + ) + log.warning( + "Tracing torchbind method %s.%s with real ScriptObject. This may" + " cause the original object being mutated. If this is not intended," + ' You can register a fake class with torch._library.register_fake_class("%s::%s").', + class_name, + method, + ns, + class_name, + ) + + ret = track_tensor_tree(out, out_proxy, constant=None, tracer=mode.tracer) + if "val" not in out_proxy.node.meta: + assert out is None or isinstance( + out, (int, float, bool) + ), "Currently, only these constant dtypes are supported to be returned from torchbind methods." + out_proxy.node.meta["val"] = out + return ret + + +# When tracing with fake script object, the call_torchbind op will return a fake tensor +# When tracing with real script object, the call_torchbind op may return a real tensor, +# we need to convert it to fake tensor mannually. Dynamic shape is surpported. +@call_torchbind.py_impl(FakeTensorMode) +def call_torchbind_fake(mode, *args, **kwargs): + with mode: + out = call_torchbind_impl(*args, **kwargs) + return pytree.tree_map_only( + torch.Tensor, + lambda x: mode.from_tensor(x, static_shapes=True) + if not isinstance(x, torch._subclasses.fake_tensor.FakeTensor) + else x, + out, + ) + + +call_torchbind.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(call_torchbind, deferred_error=True) +) + + +@call_torchbind.py_functionalize_impl +def call_torchbind_func(ctx, *args, **kwargs): + from torch._higher_order_ops.effects import handle_effects + + return handle_effects( + ctx.mode._allow_token_discovery, ctx.mode._tokens, call_torchbind, args, kwargs + ) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/triton_kernel_wrap.py b/lib/python3.12/site-packages/torch/_higher_order_ops/triton_kernel_wrap.py new file mode 100644 index 0000000000000000000000000000000000000000..ace56135fe1163ce819ca498d4f33baf71b6e1c8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/triton_kernel_wrap.py @@ -0,0 +1,1417 @@ +import collections +import copy +import dataclasses +import inspect +import logging +import threading +from collections import defaultdict +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + TYPE_CHECKING, + Union, +) +from typing_extensions import Never + +import sympy + +import torch.fx as fx +import torch.utils._pytree as pytree +from torch import SymInt, Tensor +from torch._C import DispatchKey +from torch._ops import HigherOrderOperator +from torch._prims_common import clone_preserve_strides +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ProxyTorchDispatchMode, + track_tensor_tree, +) +from torch.fx.experimental.symbolic_shapes import guard_scalar + + +if TYPE_CHECKING: + from triton._C.libtriton.ir import ( + module as TritonIRModule, + operation as TritonIROperation, + ) + + from torch._dynamo.symbolic_convert import InstructionTranslator + from torch._dynamo.variables.constant import ConstantVariable + from torch._dynamo.variables.functions import TritonKernelVariable + from torch._subclasses.functional_tensor import BaseFunctionalizeAPI + from torch.fx.proxy import Proxy + from torch.utils._triton import has_triton + + TritonMetaParamsType = Dict[str, int] + TritonGridTupleType = Tuple[Union[int, sympy.Expr, SymInt], ...] + TritonGridCallableType = Callable[[TritonMetaParamsType], Tuple[int, ...]] + TritonGridType = Union[TritonGridTupleType, TritonGridCallableType] + + if has_triton(): + from triton.runtime.autotuner import Autotuner + from triton.runtime.jit import JITFunction + else: + + class Autotuner: # type: ignore[no-redef] + pass + + class JITFunction: # type: ignore[no-redef] + pass + + TritonKernelType = Union[Autotuner, JITFunction] + + +log = logging.getLogger("torch._dynamo") + +# TMADescriptorMetadata maps kernel parameter names to the metadata that allows +# reconstructing TMA descriptors from the underlying tensors (passed as kernel +# arguments in the fx graph, instead of the TMA descriptors). Namely: a tuple +# conisting of list of dims, list of block dims, and element size. E.g., for this +# call in host-side Triton TMA API ``create_2d_tma_descriptor(ptr, 50, 60, 32, 15, 4)``, +# the metadata will look like ``([50, 60], [32, 15], 4)``. All ints can be SymInts. +TMADescriptorMetadata = Dict[ + str, # kernel parameter name + Tuple[ + List[Union[int, SymInt]], # dims + List[Union[int, SymInt]], # block_dims + Union[int, SymInt], # element_size + ], +] + + +############################################################################### +# Kernel Side Table + + +# We cannot put Triton Kernels into the FX graph as the graph nodes +# do not support arbitrary functions. +# Use a side table. +# We use two dicts so that fetching both the kernel and id are O(1) +class KernelSideTable: + id_to_kernel: Dict[int, "TritonKernelType"] = {} + kernel_to_id: Dict["TritonKernelType", int] = {} + constant_args: Dict[int, Dict[str, Any]] = {} + lock = threading.Lock() + + # Returns index on the table + def add_kernel(self, kernel: "TritonKernelType") -> int: + with self.lock: + if kernel in self.kernel_to_id: + return self.kernel_to_id[kernel] + + idx = len(self.id_to_kernel) + self.id_to_kernel[idx] = kernel + self.kernel_to_id[kernel] = idx + return idx + + # Returns the triton kernel at the given index + def get_kernel(self, idx: int) -> "TritonKernelType": + # No need to lock here as fetching from dict is atomic + assert idx in self.id_to_kernel + return self.id_to_kernel[idx] + + # Not every constant arg can be added to the graph. Use this side table + # for constant args. + def add_constant_args(self, args: Dict[str, Any]) -> int: + with self.lock: + idx = len(self.constant_args) + self.constant_args[idx] = args + return idx + + # Returns the constant args + def get_constant_args(self, idx: int) -> Dict[str, Any]: + # No need to lock here as fetching from dict is atomic + assert idx in self.constant_args + return self.constant_args[idx] + + # Resets the table (only meant to be used in unit tests) + # This is only safe assuming single threaded execution + def reset_table(self) -> None: + self.id_to_kernel = {} + self.kernel_to_id = {} + self.constant_args = {} + + +kernel_side_table = KernelSideTable() + + +############################################################################### +# Mutation Tracker + + +@dataclasses.dataclass(frozen=True) +class Param: + idx: int + + +@dataclasses.dataclass(frozen=True) +class Intermediate: + idx: int + + def fake(self) -> bool: + return self.idx < 0 + + +@dataclasses.dataclass(frozen=True) +class Op: + name: str + fn_call_name: Optional[str] + args: List[Union[Param, Intermediate]] + ret: Intermediate = dataclasses.field(repr=False) + + def __post_init__(self) -> None: + if self.name == "tt.call": + assert self.fn_call_name is not None + else: + assert self.fn_call_name is None + + +def generate_ttir( + kernel: "TritonKernelType", kwargs: Dict[str, Any] +) -> Tuple["TritonIRModule", List[str]]: + """ + Uses Triton's internal code generation to create TTIR + """ + import sympy + import triton + from triton.compiler.compiler import ASTSource + from triton.runtime.autotuner import Autotuner + from triton.runtime.jit import JITFunction + + import torch._inductor.ir + from torch._subclasses.fake_tensor import FakeTensor + + if isinstance(kernel, Autotuner): + if len(kernel.configs) > 0: + # If we are autotuning, then it doesn't matter which version gets + # picked for tracing purposes, so lets pick the first one + kwargs = {**kwargs, **kernel.configs[0].kwargs} + kernel = kernel.fn + + assert isinstance(kernel, JITFunction) + + context = triton._C.libtriton.ir.context() + target = triton.runtime.driver.active.get_current_target() + backend = triton.compiler.compiler.make_backend(target) + options = backend.parse_options({}) + + # ignore backend-specific kwargs same way as in the native Triton code + # https://github.com/triton-lang/triton/blob/a6bb57d6285e723c58e87dd7cba263db6efff789/python/triton/runtime/jit.py#L594-L596 + # why this is important for user-defined Triton kernels on AMD: https://github.com/pytorch/pytorch/issues/140800 + for name in list(kwargs): + if name not in kernel.arg_names and name in options.__dict__: + kwargs.pop(name) + + if len(kwargs) != len(kernel.arg_names): + raise ValueError( + "Incorrect number of arguments passed to kernel: " + f"passed {list(kwargs.keys())}, expected {kernel.arg_names}." + ) + + # Replace all SymExprs with a regular value for TTIR generation + # Replace all FakeTensor/TensorBox with real tensors + # These replacements are needed for triton's type, key and config functions + ordered_args: Dict[str, Any] = {} + for name in kernel.arg_names: + a = kwargs[name] + if isinstance(a, (torch.SymInt, torch.SymFloat, torch.SymBool, sympy.Expr)): + ordered_args[name] = 2 + elif isinstance(a, (FakeTensor, torch._inductor.ir.TensorBox)): + with torch._C._DisableTorchDispatch(): + ordered_args[name] = torch.empty(2, dtype=a.dtype) + else: + ordered_args[name] = a + + ordered_tensor_names = [ + name for name, arg in ordered_args.items() if isinstance(arg, Tensor) + ] + + def _get_specialization(args): # type: ignore[no-untyped-def] + try: + from triton.backends.compiler import AttrsDescriptor # noqa: F401 + + target = triton.runtime.driver.active.get_current_target() + backend = triton.compiler.compiler.make_backend(target) + return backend.get_attrs_descriptor(args, kernel.params) + except ImportError: + return kernel._get_config(*args) + + specialization = _get_specialization(ordered_args.values()) + constants = { + name: arg for name, arg in ordered_args.items() if not isinstance(arg, Tensor) + } + + # Build kernel signature -- doesn't include constexpr arguments. + signature = { + name: kernel._type_of(kernel._key_of(arg)) + for i, (name, arg) in enumerate(ordered_args.items()) + if i not in kernel.constexprs + } + + triton._C.libtriton.ir.load_dialects(context) + backend.load_dialects(context) + + src = ASTSource(kernel, signature, constants, specialization) + + # Triton changes ASTSource.make_ir to take 3/4 arguments. Handle + # backward compatibility here. + make_ir_sig_params = len(inspect.signature(src.make_ir).parameters) + if make_ir_sig_params == 2: + ttir_module = src.make_ir(options, context) + elif make_ir_sig_params == 3: + codegen_fns = backend.get_codegen_implementation() + ttir_module = src.make_ir(options, codegen_fns, context) + else: + codegen_fns = backend.get_codegen_implementation() + module_map = backend.get_module_map() + ttir_module = src.make_ir(options, codegen_fns, module_map, context) + if not ttir_module.verify(): + raise RuntimeError("Verification for TTIR module has failed") + + return ttir_module, ordered_tensor_names + + +def ttir_to_functions( + ttir_module: "TritonIRModule", +) -> Dict[str, Dict[Intermediate, List[Op]]]: + """ + Walk the `ttir_module` bottom up to mine the `functions` from + the structured MLIR entities representing the Triton kernel + (mlir::Operation, mlir::Block, mlir::Region). + """ + functions: Dict[str, Dict[Intermediate, List[Op]]] = {} + + # block id --> op result (Intermediate) --> one or more ops + op_stack: Dict[int, Dict[Intermediate, List[Op]]] = defaultdict( + lambda: defaultdict(list) + ) + region_id_to_block_ids: Dict[int, List[int]] = defaultdict(list) + block_id_to_block_arg_ids: Dict[int, List[int]] = {} + replacements: Dict[int, Union[Intermediate, Param]] = {} + reindex_map: Dict[int, int] = {} + next_fake_intermediate = 0 + + def reindex(idx: int) -> int: + if idx not in reindex_map: + reindex_map[idx] = len(reindex_map) + return reindex_map[idx] + + def mlir_to_functions(op: "TritonIROperation") -> None: + name: str = op.get_name() + if name == "builtin.module": + # this wraps all tt.func ops + return + + operand_ids: List[int] = [ + reindex(op.get_operand(i).id()) for i in range(op.get_num_operands()) + ] + result_ids: List[int] = [ + reindex(op.get_result(i).id()) for i in range(op.get_num_results()) + ] + + child_block_ids: List[int] = [] + for i in [op.get_region(i).id() for i in range(op.get_num_regions())]: + # as the walk is bottom-up, the region_id_to_block_ids[i] + # must be populated by the time we process the enclosing op + child_block_ids.extend(region_id_to_block_ids[i]) + + parent_block_id = -1 + parent_block = op.get_block() + if parent_block is not None: + parent_block_id = parent_block.id() + if parent_block_id not in block_id_to_block_arg_ids: + block_id_to_block_arg_ids[parent_block_id] = [] + for i in range(parent_block.get_num_arguments()): + block_id_to_block_arg_ids[parent_block_id].append( + reindex(parent_block.get_argument(i).id()), + ) + # the region info is collected via ops' parent blocks to be + # used later when the region's encloding op is traversed + parent_region = parent_block.get_parent() + if parent_region is not None: + region_id_to_block_ids[parent_region.id()].append(parent_block_id) + + nonlocal next_fake_intermediate + + if name == "tt.func": + # for function ops: gather and inline + # the ops from all child blocks + fn_ops = defaultdict(list) + for child_block_id in child_block_ids: + for result, block_fn_ops in op_stack.pop(child_block_id).items(): + for block_fn_op in block_fn_ops: + fn_ops[result].append(block_fn_op) + + # replace the corresponding Intermediates in the + # child op args with the function args (Params) + for i, idx in enumerate(block_id_to_block_arg_ids[child_block_ids[0]]): + replacements[idx] = Param(i) + + for fn_op_list in fn_ops.values(): + for fn_op in fn_op_list: + for i in range(len(fn_op.args)): + arg = fn_op.args[i] + seen = set() # to break cycles + # there can be transitive replacements, but likely + # no cycles (we keep the `seen` set just in case) + while ( + isinstance(arg, Intermediate) + and arg.idx in replacements + and arg.idx not in seen + ): + seen.add(arg.idx) + arg = fn_op.args[i] = replacements[arg.idx] + + # next function capture starts + # with empty replacements + replacements.clear() + + fn_name = op.get_str_attr("sym_name") + functions[fn_name] = fn_ops + elif child_block_ids: + if name in {"scf.if", "scf.for", "scf.while", "tt.reduce", "tt.scan"}: + # for blocked ops: inline the enclosed ops into + # the parent block + rewire the last op in each + # child block to return the block result + return_ops = [] + for block_id in child_block_ids: + if name == "scf.for": + # example: + # %result = scf.for %iv = %lb to %ub step %step iter_args(%arg = %init) -> (i32) ... + # block args: 2 (%iv, %arg) + # op operands: 4 (%lb, %ub, %step, %init) + # `%arg` is mapping to `%init` + for i, idx in enumerate(block_id_to_block_arg_ids[block_id]): + if i == 0: + next_fake_intermediate -= 1 + replacements[idx] = Intermediate(next_fake_intermediate) + else: + replacements[idx] = Intermediate(operand_ids[i + 2]) + elif name == "scf.while": + # example: + # %3:3 = scf.while (%arg2 = %1, %arg3 = %2, %arg4 = %c0_i32_8) ... + # block args: 3 (%arg2, %arg3, %arg4) + # op operands: 3 (%1, %2, %c0_i32_8) + # `%arg2` is mapping to `%1`, `%arg3` is mapping to `%2`, ... + for i, idx in enumerate(block_id_to_block_arg_ids[block_id]): + replacements[idx] = Intermediate(operand_ids[i]) + elif name == "scf.if": + # the scf block args are ignored by the pass. but, as they + # may be used as operands of the ops inside the block + # (and nested blocks inlined in the current block by now), + # they are replaced by new fake Intermediates to avoid "this + # operand is not returned by any other op in the fn" error + # in the downstream analysis + for idx in block_id_to_block_arg_ids[block_id]: + next_fake_intermediate -= 1 + replacements[idx] = Intermediate(next_fake_intermediate) + else: + assert name in ("tt.reduce", "tt.scan") + # wire the block arguments to the op arguments + num_operands = len(operand_ids) + block_arg_ids = block_id_to_block_arg_ids[block_id] + assert len(block_arg_ids) == 2 * num_operands, ( + f"{name} is expected to have twice as " + "many block arguments as op arguments: " + f"{operand_ids=}, {block_arg_ids=}." + ) + for i, idx in enumerate(block_arg_ids): + # for a tt.reduce/tt.scan op with N arguments, the block + # arguments comprise N reduced values followed by + # N current values corresponding to the N op args + replacements[idx] = Intermediate( + operand_ids[i % num_operands] + ) + + if block_id in op_stack: + block_ops = op_stack.pop(block_id) + if not block_ops: + continue + last_ret, last_ops = block_ops.popitem() + if all( + op.name + in ("scf.yield", "tt.reduce.return", "tt.scan.return") + for op in last_ops + ): + # if last_ops are all return ops, treat them separately + return_ops.extend(last_ops) + else: + # otherwise, return last_ops to the block + block_ops[last_ret] = last_ops + for op_result, child_ops in block_ops.items(): + op_stack[parent_block_id][op_result].extend(child_ops) + + scf_results = [Intermediate(idx) for idx in result_ids] + for scf_result in scf_results: + for return_op in return_ops: + op_stack[parent_block_id][scf_result].append(return_op) + else: + raise RuntimeError( + f"Unknown blocked function: {name}. Can't capture the TTIR." + ) + else: + callee = None + if name == "tt.call": + callee = op.get_flat_symbol_ref_attr("callee") + args: List[Union[Param, Intermediate]] = [ + Intermediate(operand) for operand in operand_ids + ] + block_ops = op_stack[parent_block_id] + if result_ids: + for result_id in result_ids: + res = Intermediate(result_id) + block_ops[res].append(Op(name, callee, args, res)) + else: + next_fake_intermediate -= 1 + fake_res = Intermediate(next_fake_intermediate) + block_ops[fake_res].append(Op(name, callee, args, fake_res)) + + ttir_module.walk(mlir_to_functions) + + return functions + + +class MemoizeWithCycleCheck: + fn: Callable[..., Any] + cache: Dict[Tuple[str, int], Any] + + def __init__(self, fn: Callable[..., Any]) -> None: + self.fn = fn + self.reset() + + def __call__( + self, + functions: Dict[str, Dict[Intermediate, List[Op]]], + fn_name: str, + num_args: int, + ) -> List[bool]: + key = (fn_name, num_args) + if key not in self.cache: + self.cache[key] = None + self.cache[key] = self.fn(functions, fn_name, num_args) + if self.cache[key] is None: + raise RuntimeError("Recursion is not supported") + return self.cache[key] + + def reset(self) -> None: + self.cache = {} + + +@MemoizeWithCycleCheck +def analyze_kernel_mutations( + functions: Dict[str, Dict[Intermediate, List[Op]]], fn_name: str, num_args: int +) -> List[bool]: + """ + Analyzes the graph to detect all sinks from a predefined list of sinks + by using triton's MemWrite trait list. NOTE: What if triton exposed this? + From each sink, it traverses the CFG backwards to identify all the input + pointers that are mutated. + """ + # Name of mutation op to mutated parameter indices + # List from Triton Github include/triton/Dialect/Triton/IR/TritonOps.td + # All the OPs that have MemWrite trait. + # What if Triton exposed this? + MUTATION_OPS = { + "tt.store": [0], + "tt.atomic_cas": [0], + "tt.atomic_rmw": [0], + "tt.experimental_descriptor_store": [0], + } + # Ops that we want to bail out on + UNKNOWN_OPS = {"tt.elementwise_inline_asm"} + + stack: List[Union[Param, Intermediate]] = [] + visited = set() + ops = functions[fn_name] + for op_list in ops.values(): + for op in op_list: + if op.name in UNKNOWN_OPS: + raise RuntimeError( + f"ttir analysis hit an op we do not know how to analyze: {op.name}" + ) + + if op.name == "tt.call": + assert op.fn_call_name in functions + mutations = analyze_kernel_mutations( + functions, op.fn_call_name, len(op.args) + ) + stack.extend(arg for arg, mutated in zip(op.args, mutations) if mutated) + else: + stack.extend(op.args[idx] for idx in MUTATION_OPS.get(op.name, [])) + + # The following is an iterative DFS algorithm + mutated = [False] * num_args + while stack: + arg = stack.pop() + if arg in visited: + continue + + visited.add(arg) + + if isinstance(arg, Param): + if arg.idx >= num_args: + # This is an argument defined in the kernel, not passed in + continue + mutated[arg.idx] = True + elif isinstance(arg, Intermediate) and not arg.fake(): + for op in ops[arg]: + # Skip arguments to load + if op.name != "tt.load": + stack.extend(op.args) + return mutated + + +def identify_mutated_tensors( + kernel: "TritonKernelType", kwargs: Dict[str, Any] +) -> List[str]: + """ + Given a triton kernel and the arguments for this kernel, this function + 1) Retrieves the TTIR converted version of the kernel from Triton's API. + 2) Parses the TTIR and creates a control flow graph + 3) Analyzes the graph to detect all input tensor mutations + """ + + ttir_module = None + functions = None + try: + ttir_module, ordered_tensor_names = generate_ttir(kernel, kwargs) + + # extract functions from TTIR using MLIR bindings exposed by Triton code + functions = ttir_to_functions(ttir_module) + + assert functions is not None + kernel_name = next(iter(functions.keys())) + # Triton codegen modifies the name + assert kernel.fn.__name__ in kernel_name + # Reset the cache between top level invocations + # The cache for analyze kernel mutations is mainly used for cycle + # detection, so each top level invocation needs a clean cache + analyze_kernel_mutations.reset() + mutations = analyze_kernel_mutations( + functions, kernel_name, len(ordered_tensor_names) + ) + + return [ + ordered_tensor_names[i] for i, mutated in enumerate(mutations) if mutated + ] + except Exception as e: + log.warning( + "Encountered an exception in identify_mutated_tensors, assuming every input is mutated", + exc_info=True, + ) + if ttir_module is not None: + log.debug("TTIR:\n%s", str(ttir_module)) + if functions is not None: + log.debug("functions:") + for name, fn in functions.items(): + log.debug("===\t%s\t===", name) + for ret, ops in fn.items(): + log.debug("%s\t=>\t%s", ret, ops) + return [key for key, value in kwargs.items() if isinstance(value, Tensor)] + + +############################################################################### +# Triton Kernel Wrappers + + +# Used for wrapping a Triton Kernel +class TritonKernelWrapperMutation(HigherOrderOperator): + def __init__(self) -> None: + super().__init__("triton_kernel_wrapper_mutation", cacheable=True) + + def __call__( + self, + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], + ) -> Any: + return super().__call__( + kernel_idx=kernel_idx, + constant_args_idx=constant_args_idx, + grid=grid, + tma_descriptor_metadata=tma_descriptor_metadata, + kwargs=kwargs, + ) + + +triton_kernel_wrapper_mutation = TritonKernelWrapperMutation() + + +# Used for wrapping a Triton Kernel in a functional manner +class TritonKernelWrapperFunctional(HigherOrderOperator): + def __init__(self) -> None: + super().__init__("triton_kernel_wrapper_functional", cacheable=True) + + def __call__( + self, + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], + tensors_to_clone: List[str], + ) -> Dict[str, Any]: + return super().__call__( + kernel_idx=kernel_idx, + constant_args_idx=constant_args_idx, + grid=grid, + tma_descriptor_metadata=tma_descriptor_metadata, + kwargs=kwargs, + tensors_to_clone=tensors_to_clone, + ) + + +triton_kernel_wrapper_functional = TritonKernelWrapperFunctional() + + +@triton_kernel_wrapper_mutation.py_impl(DispatchKey.CompositeExplicitAutograd) +def triton_kernel_wrapper_mutation_dense( + *, + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], +) -> None: + from torch._inductor.codegen.wrapper import user_defined_kernel_grid_fn_code + + kernel = kernel_side_table.get_kernel(kernel_idx) + constant_args = kernel_side_table.get_constant_args(constant_args_idx) + + if len(grid) == 1: + grid_fn = grid[0] + else: + fn_name, code = user_defined_kernel_grid_fn_code( + kernel.fn.__name__, kernel.configs, grid + ) + namespace: Dict[str, Any] = {} + exec(code, namespace) + grid_fn = namespace[fn_name] + + if tma_descriptor_metadata: + from triton.tools.experimental_descriptor import ( # noqa: F401 + create_1d_tma_descriptor, + create_2d_tma_descriptor, + ) + + # as we need to launch the kernel here, we "unwrap" the + # tma_descriptor_metadata, create the TMA descriptors + # from it, and replace the tensors in the kwargs by the + # correspoinding TMA descriptors before launching + kwargs = kwargs.copy() + for k, v in tma_descriptor_metadata.items(): + tensor = kwargs[k] + dims, block_dims, element_size = v + create_tma_descriptor = ( + create_1d_tma_descriptor if len(dims) == 1 else create_2d_tma_descriptor + ) + kwargs[k] = create_tma_descriptor( + tensor.data_ptr(), + *dims, + *block_dims, + element_size, + ) + + # move as many positional arguments from dicts to args as we + # can to circumvent the bug with the kwargs and pre_/post_hook: + # https://github.com/triton-lang/triton/issues/5082 + # TODO: remove this when the Triton issue above is fixed + args = [] + # copy kwargs and constant_args here to + # avoid mutating the original inputs + kwargs = kwargs.copy() + constant_args = constant_args.copy() + for name in kernel.arg_names: + if name in kwargs: + args.append(kwargs.pop(name)) + elif name in constant_args: + args.append(constant_args.pop(name)) + else: + break + + kernel[grid_fn](*args, **kwargs, **constant_args) + + +@triton_kernel_wrapper_mutation.py_impl(FakeTensorMode) +def triton_kernel_wrapper_mutation_fake_tensor_mode( + mode: FakeTensorMode, + *, + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], +) -> None: + with mode: + return None + + +@triton_kernel_wrapper_mutation.py_impl(DispatchKey.Meta) +def _( + *, + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], +) -> None: + return None + + +def trace_triton_kernel_wrapper( + proxy_mode: ProxyTorchDispatchMode, + func_overload: Callable[..., Any], + node_args: Dict[str, Any], +) -> Optional[Dict[str, Any]]: + with disable_proxy_modes_tracing(): + out = func_overload(**node_args) + + proxy_args = pytree.tree_map( + proxy_mode.tracer.unwrap_proxy, node_args # type: ignore[union-attr] + ) + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", + func_overload, + (), + proxy_args, + name=func_overload.__name__ + "_proxy", + ) + + ret = track_tensor_tree(out, out_proxy, constant=None, tracer=proxy_mode.tracer) + return ret + + +@triton_kernel_wrapper_mutation.py_impl(ProxyTorchDispatchMode) +def triton_kernel_wrapper_mutation_proxy_torch_dispatch_mode( + mode: ProxyTorchDispatchMode, + *, + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], +) -> None: + trace_triton_kernel_wrapper( + mode, + triton_kernel_wrapper_mutation, + { + "kernel_idx": kernel_idx, + "constant_args_idx": constant_args_idx, + "grid": grid, + "tma_descriptor_metadata": tma_descriptor_metadata, + "kwargs": kwargs, + }, + ) + + return None + + +def get_mutated_tensors( + kernel_idx: int, constant_args_idx: int, kwargs: Dict[str, Any] +) -> List[str]: + kernel = kernel_side_table.get_kernel(kernel_idx) + constant_args = kernel_side_table.get_constant_args(constant_args_idx) + return identify_mutated_tensors(kernel, {**kwargs, **constant_args}) + + +@triton_kernel_wrapper_mutation.py_functionalize_impl +def triton_kernel_wrapper_mutation_functionalize( + ctx: "BaseFunctionalizeAPI", + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], +) -> None: + unwrapped_kwargs = ctx.unwrap_tensors(kwargs) # type: ignore[arg-type] + # TODO(oulgen): Preexisting bug, if two kernel inputs are views of each + # other, and one gets mutated in kernel, and later another gets mutated, + # they are no longer equal. Fix this by graph breaking on this condition + # earlier in dynamo. + tensors_to_clone = get_mutated_tensors( + kernel_idx, constant_args_idx, unwrapped_kwargs + ) + with ctx.redispatch_to_next(): + unwrapped_outputs = triton_kernel_wrapper_functional( + kernel_idx=kernel_idx, + constant_args_idx=constant_args_idx, + grid=grid, + tma_descriptor_metadata=tma_descriptor_metadata, + kwargs=unwrapped_kwargs, + tensors_to_clone=tensors_to_clone, + ) + + assert set(unwrapped_outputs.keys()).issubset(set(kwargs.keys())) + for key, output_arg in unwrapped_outputs.items(): + if not isinstance(output_arg, Tensor): + continue + input_arg = kwargs[key] + assert isinstance(input_arg, Tensor) + + ctx.replace(input_arg, output_arg) + # indicate that above replace is hidden from autograd + ctx.mark_mutation_hidden_from_autograd(input_arg) + ctx.commit_update(input_arg) + ctx.sync(input_arg) + return None + + +@triton_kernel_wrapper_functional.py_impl(DispatchKey.CompositeExplicitAutograd) +def triton_kernel_wrapper_functional_dense( + *, + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], + tensors_to_clone: List[str], +) -> Dict[str, Any]: + # TODO(oulgen): For performance reasons, we want to ensure that these + # `clone_preserve_strides` calls are never executed at runtime + # (inductor should always optimize them away). + # Requires https://github.com/pytorch/pytorch/issues/109240 + kwargs = { + key: (clone_preserve_strides(val) if key in tensors_to_clone else val) + for key, val in kwargs.items() + } + triton_kernel_wrapper_mutation( + kernel_idx=kernel_idx, + constant_args_idx=constant_args_idx, + grid=grid, + tma_descriptor_metadata=tma_descriptor_metadata, + kwargs=kwargs, + ) + return {key: val for key, val in kwargs.items() if key in tensors_to_clone} + + +@triton_kernel_wrapper_functional.py_impl(FakeTensorMode) +def triton_kernel_wrapper_functional_fake_tensor_mode( + mode: FakeTensorMode, + *, + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], + tensors_to_clone: List[str], +) -> Dict[str, Any]: + # TODO(oulgen): For performance reasons, we want to ensure that these + # `clone_preserve_strides` calls are never executed at runtime + # (inductor should always optimize them away). + # Requires https://github.com/pytorch/pytorch/issues/109240 + with mode: + return { + key: clone_preserve_strides(val) + for key, val in kwargs.items() + if key in tensors_to_clone + } + + +@triton_kernel_wrapper_functional.py_impl(ProxyTorchDispatchMode) +def triton_kernel_wrapper_functional_proxy_torch_dispatch_mode( + mode: ProxyTorchDispatchMode, + *, + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], + tensors_to_clone: List[str], +) -> Dict[str, Any]: + ret = trace_triton_kernel_wrapper( + mode, + triton_kernel_wrapper_functional, + { + "kernel_idx": kernel_idx, + "constant_args_idx": constant_args_idx, + "grid": grid, + "tma_descriptor_metadata": tma_descriptor_metadata, + "kwargs": kwargs, + "tensors_to_clone": tensors_to_clone, + }, + ) + assert ret is not None + return ret + + +@triton_kernel_wrapper_functional.py_functionalize_impl +def triton_kernel_wrapper_functional_functionalize( + ctx: "BaseFunctionalizeAPI", + kernel_idx: int, + constant_args_idx: int, + grid: List["TritonGridType"], + tma_descriptor_metadata: TMADescriptorMetadata, + kwargs: Dict[str, Any], + tensors_to_clone: List[str], +) -> Dict[str, Any]: + unwrapped_kwargs = ctx.unwrap_tensors(kwargs) # type: ignore[arg-type] + with ctx.redispatch_to_next(): + outputs = triton_kernel_wrapper_functional( + kernel_idx=kernel_idx, + constant_args_idx=constant_args_idx, + grid=grid, + tma_descriptor_metadata=tma_descriptor_metadata, + kwargs=unwrapped_kwargs, + tensors_to_clone=tensors_to_clone, + ) + return ctx.wrap_tensors(outputs) # type: ignore[return-value,arg-type] + + +triton_kernel_wrapper_mutation.fallthrough(DispatchKey.PythonDispatcher) # type: ignore[attr-defined] +triton_kernel_wrapper_mutation.fallthrough(DispatchKey.PythonTLSSnapshot) # type: ignore[attr-defined] +triton_kernel_wrapper_mutation.fallthrough(DispatchKey.ADInplaceOrView) +triton_kernel_wrapper_mutation.fallthrough(DispatchKey.BackendSelect) +triton_kernel_wrapper_mutation.fallthrough(DispatchKey.AutocastCPU) # type: ignore[attr-defined] +triton_kernel_wrapper_mutation.fallthrough(DispatchKey.AutocastCUDA) # type: ignore[attr-defined] +triton_kernel_wrapper_mutation.fallthrough(DispatchKey.AutogradCUDA) +triton_kernel_wrapper_mutation.fallthrough(DispatchKey.AutogradCPU) + +triton_kernel_wrapper_functional.fallthrough(DispatchKey.PythonDispatcher) # type: ignore[attr-defined] +triton_kernel_wrapper_functional.fallthrough(DispatchKey.PythonTLSSnapshot) # type: ignore[attr-defined] +triton_kernel_wrapper_functional.fallthrough(DispatchKey.ADInplaceOrView) +triton_kernel_wrapper_functional.fallthrough(DispatchKey.BackendSelect) +triton_kernel_wrapper_functional.fallthrough(DispatchKey.AutocastCPU) # type: ignore[attr-defined] +triton_kernel_wrapper_functional.fallthrough(DispatchKey.AutocastCUDA) # type: ignore[attr-defined] +triton_kernel_wrapper_functional.fallthrough(DispatchKey.AutogradCUDA) +triton_kernel_wrapper_functional.fallthrough(DispatchKey.AutogradCUDA) +triton_kernel_wrapper_functional.fallthrough(DispatchKey.AutogradCPU) + + +############################################################################### +# The "TritonHOPifier": a class that transforms a call to a triton kernel into +# a call to the triton_kernel_wrapper_mutation HOP. + + +class TritonHOPifier: + """Orchestrator for converting a user-defined triton kernel into a call + to the triton_kernel_wrapper_mutation HOP. + + It has two main use cases. + + 1. When Dynamo sees a triton kernel, it wraps it into a TritonKernelVariable + and uses the TritonHOPifier to convert calls to the TritonKernelVariable + into a call to the HOP. + + 2. In order to capture a user-defined triton kernel while performing + tracing (via make_fx or non-strict export), a user must annotate their + triton kernel with the `wrap_triton` decorator. The decorator uses + TritonHOPifier to convert calls to the triton kernel into a call + to the HOP (which can then be traced). + + Because Dynamo has its own calling conventions for e.g. invoking a user-defined function + TritonHOPifier is an abstract class that can be overriden by its subclasses. + """ + + def raise_unsupported(self, msg: str) -> Never: + raise NotImplementedError("abstract method") + + def is_callable(self, maybe_callable: Any) -> bool: + raise NotImplementedError("abstract method") + + def get_value(self, val: Any) -> Any: + raise NotImplementedError("abstract method") + + def call_grid( # type: ignore[no-untyped-def] + self, + grid, + meta, + tx, + ) -> Union[Tuple[Union[int, sympy.Expr, SymInt], ...], Tuple["Proxy", ...]]: + raise NotImplementedError("abstract method") + + def call_HOP( # type: ignore[no-untyped-def] + self, + variable, + grids, + combined_args: Dict[str, Any], + tx, + ) -> Optional["ConstantVariable"]: + raise NotImplementedError("abstract method") + + def check_grid( # type: ignore[no-untyped-def] + self, grid + ) -> Union[Tuple[Union[int, sympy.Expr, SymInt], ...], Tuple["Proxy", ...]]: + raise NotImplementedError("abstract method") + + def init_variable( + self, + variable: Union["TraceableTritonKernelWrapper", "TritonKernelVariable"], + kernel: "TritonKernelType", + kernel_idx: Optional[int], + grid: Optional["TritonGridType"], + ) -> None: + from triton.runtime.autotuner import Autotuner + + assert kernel is not None + + variable.kernel = kernel + variable.kernel_idx = kernel_side_table.add_kernel(kernel) + + assert kernel_idx is None or variable.kernel_idx == kernel_idx + + variable.grid = grid + + if isinstance(kernel, Autotuner): + import torch + import torch._dynamo + + # We only support configs, keys, and restore_value arguments + # of triton.autotune. Make sure other arguments are defaulted. + defaults = inspect.signature(Autotuner.__init__).parameters + # Newer version of triton change attribute name from warmup to num_warmup and rep to num_rep. + # The call to get_first_attr is to maintain backward-compatibility. + if ( + not torch._inductor.config.unsafe_ignore_unsupported_triton_autotune_args + and ( + ( + "warmup" in defaults + and defaults["warmup"].default + != torch._dynamo.utils.get_first_attr( + kernel, "num_warmups", "warmup" + ) + ) + or ( + "rep" in defaults + and defaults["rep"].default + != torch._dynamo.utils.get_first_attr(kernel, "num_reps", "rep") + ) + or ( + "prune_configs_by" in defaults + and defaults["prune_configs_by"].default + != kernel.early_config_prune + ) + or ( + "use_cuda_graph" in defaults + and defaults["use_cuda_graph"].default != kernel.use_cuda_graph + ) + ) + ): + self.raise_unsupported( + "Only configs, keys, restore_value, and reset_to_zero are supported for triton.autotune" + ) + if ( + not torch._inductor.config.unsafe_ignore_unsupported_triton_autotune_args + and ( + # pre_hook requires running arbitrary code at runtime, which we cannot handle at this time + # https://github.com/pytorch/pytorch/issues/139059 + # we can't support pre_hook or post_hook in user defined triton kernels at the moment, + # as they require the ability to execute code at runtime (AOTI can't support this) + ( + hasattr(kernel, "user_defined_pre_hook") + and kernel.user_defined_pre_hook is not False + ) + or ( + hasattr(kernel, "user_defined_post_hook") + and kernel.user_defined_post_hook is not False + ) + or ( + # Check Config passed to autotuner in configs + any(cfg.pre_hook is not None for cfg in kernel.configs) + ) + ) + ): + self.raise_unsupported( + "pre_hook and post_hook are not supported in triton.Autotune or triton.Config" + ) + + def call_getitem( + self, + variable: Union["TritonKernelVariable", "TraceableTritonKernelWrapper"], + args: Sequence[Any], + ) -> Union["TritonKernelVariable", "TraceableTritonKernelWrapper"]: + # __getitem__ should only be called if we don't already have a grid + # Only grid needs to be passed + if variable.grid is not None or len(args) != 1: + self.raise_unsupported( + "Triton kernels should be called with only a single grid" + ) + + return type(variable)( + kernel=variable.kernel, + kernel_idx=variable.kernel_idx, + grid=args[0], + ) + + def call_run( + self, + variable: Union["TritonKernelVariable", "TraceableTritonKernelWrapper"], + args: Sequence[Any], + kwargs: Dict[str, Any], + tx: Optional["InstructionTranslator"], + ) -> Optional["ConstantVariable"]: + if "grid" not in kwargs: + self.raise_unsupported("Triton kernel requires to be called with a grid") + grid = kwargs.pop("grid") + kwargs.pop("warmup", None) + # rewrite kernel.run(*args, grid=grid) to kernel[grid](*args) + return self.call_triton_kernel( + type(variable)( + kernel=variable.kernel, kernel_idx=variable.kernel_idx, grid=grid + ), + args, + kwargs, + tx, + ) + + def call_triton_kernel( + self, + variable: Union["TritonKernelVariable", "TraceableTritonKernelWrapper"], + args: Sequence[Any], + kwargs: Dict[str, Any], + tx: Optional["InstructionTranslator"], + ) -> Optional["ConstantVariable"]: + from triton import JITFunction + from triton.runtime.autotuner import autotune, Autotuner, Config + + SPECIAL_CONFIG_NAMES = {"num_warps", "num_stages", "num_ctas"} + + if "num_ctas" in kwargs: + self.raise_unsupported( + "Passing num_ctas directly to the Triton kernel is not supported. " + "Please use a Config in @triton.autotune instead." + ) + + special_kwargs = {} + for name in SPECIAL_CONFIG_NAMES: + if name in kwargs: + # remove special kwargs from `kwargs` + val = kwargs.pop(name) + special_kwargs[name] = self.get_value(val) + + if special_kwargs: + if isinstance(variable.kernel, Autotuner): + # if there is Autotuner already, set + # special kwargs to each of its configs + new_configs = copy.deepcopy(variable.kernel.configs) + for config in new_configs: + config.__dict__.update(special_kwargs) + new_kernel = autotune(configs=new_configs, key=[])(variable.kernel.fn) + else: + # if there is no Autotuner, wrap the kernel into a + # new one with a single config with special kwargs + new_config = Config(kwargs={}, **special_kwargs) + new_kernel = autotune(configs=[new_config], key=[])(variable.kernel) + + # create a new variable to contain the new (wrapped) kernel; + # skip kernel_idx to get a new record in the kernel side table + new_var = type(variable)(new_kernel, None, variable.grid) + return self.call_triton_kernel(new_var, args, kwargs, tx) + + if isinstance(variable.kernel, Autotuner): + special_param_names = [] + for name in SPECIAL_CONFIG_NAMES: + if name in variable.kernel.fn.arg_names: + special_param_names.append(name) + + if special_param_names: + # If the Triton kernel has SPECIAL_CONFIG_NAMES in parameters, those should + # be passed from the kernel configs: the behavior of Triton runtime is that + # those values get folded into the kernel arguments iff there are parameters + # with the same name. Normally the values of those parameters are defined + # outside the `kwargs` part of the autotuning configs. Here we move them to + # the `kwargs` part (if they're absent there) to facilitate passing them as + # arguments to the kernel downstream. + updated = False + new_configs = copy.deepcopy(variable.kernel.configs) + for config in new_configs: + for name in special_param_names: + if name not in config.__dict__["kwargs"]: + assert ( + name in config.__dict__ + ), f"{name} must be in autotuning configs to be used as a kernel parameter" + config.__dict__["kwargs"][name] = config.__dict__[name] + updated = True + + if updated: + new_kernel = autotune(configs=new_configs, key=[])( + variable.kernel.fn + ) + new_var = type(variable)(new_kernel, None, variable.grid) + return self.call_triton_kernel(new_var, args, kwargs, tx) + + if variable.grid is None: + self.raise_unsupported("Triton kernels should always be called with a grid") + + # Both for grid's meta as well as for the kernel, we need combined + # args and kwargs combined and normalized + combined_args_raw = {**dict(zip(variable.kernel.arg_names, args)), **kwargs} + + configs = ( + [config.kwargs for config in variable.kernel.configs] + if isinstance(variable.kernel, Autotuner) + else [{}] + ) + grids = [] + for config_args in configs: + # If the grid is a function, then lets execute it and convert it to + # a list + grid = variable.grid + assert grid is not None + if self.is_callable(grid): + # Populate the special "meta" argument to call the grid function + meta = {**combined_args_raw, **config_args} + grid = self.call_grid(grid, meta, tx) # type: ignore[arg-type] + grids.append(self.check_grid(grid)) + + for i in range(len(grids)): + if not isinstance(grids[i], tuple): + self.raise_unsupported("Only tuple grids are supported") + # inductor expects all grids to be 3-tuple so lets make it + if len(grids[i]) == 1: + grids[i] = (grids[i][0], 1, 1) + elif len(grids[i]) == 2: + grids[i] = (grids[i][0], grids[i][1], 1) + elif len(grids[i]) > 3: + self.raise_unsupported("Grid can have at most rank 3") + + assert len(grids) != 0 + if isinstance(variable.kernel, JITFunction): + constexprs = variable.kernel.constexprs + else: + assert isinstance(variable.kernel, Autotuner) + constexprs = variable.kernel.fn.constexprs + + for idx, arg_name in enumerate(variable.kernel.arg_names): + if idx in constexprs: + if arg_name in combined_args_raw: + # [Note: Specialize tl.constexpr args in user-defined triton kernels] + # This arg is marked as tl.constexpr. That means that triton will recompile every time + # this value changes. + # https://github.com/pytorch/pytorch/issues/136504 + # One option is to correctly pass the symints in so that the symbolic expressions are defined + # when the triton code is being executed. + # But since triton will have to recompile either way, we instead just specialize on the value. + # + # Depending on the type of `variable` we might expect different types for the symbolic args: + # either SymNodeVariables (for TritonKernelVariables) or SymInts (TracingTritonKernelWrapper) + combined_args_raw[arg_name] = variable.specialize_symbolic( + combined_args_raw[arg_name] + ) + return self.call_HOP(variable, grids, combined_args_raw, tx) + + +############################################################################### +# Helpers for wrap_triton API that makes a user-defined triton kernel traceable into +# a graph via make_fx or non-strict export (coming soon) + + +class TracingTritonHOPifier(TritonHOPifier): + def raise_unsupported(self, msg: str) -> Never: + raise RuntimeError(msg) + + def is_callable(self, maybe_callable: Any) -> bool: + return callable(maybe_callable) + + def get_value(self, val: Any) -> Any: + return val + + def call_grid( + self, + grid: "TritonGridCallableType", + meta: "TritonMetaParamsType", + tx: None, + ) -> Tuple[Union[int, sympy.Expr, SymInt], ...]: + assert tx is None + assert isinstance(meta, dict) + assert callable(grid) + return grid(meta) + + def check_grid( + self, + grid: "TritonGridType", + ) -> Tuple[Union[int, sympy.Expr, SymInt], ...]: + if not isinstance(grid, collections.abc.Sequence): + raise RuntimeError( + "wrap_triton can only handle grids that resolve to Sequence[int]." + ) + # normalize to tuple + return tuple(grid) + + def call_HOP( + self, + variable: "TraceableTritonKernelWrapper", + grids: List["TritonGridTupleType"], + combined_args: Dict[str, Any], + tx: None, + ) -> None: + assert tx is None + assert isinstance(variable, TraceableTritonKernelWrapper) + + def is_graphable(val: Any) -> bool: + return isinstance(val, fx.node.base_types) + + non_graphable_args = { + k: v for k, v in combined_args.items() if not is_graphable(v) + } + graphable_args = {k: v for k, v in combined_args.items() if is_graphable(v)} + + constant_args_idx = kernel_side_table.add_constant_args(non_graphable_args) + assert isinstance(variable.kernel_idx, int) + return triton_kernel_wrapper_mutation( + kernel_idx=variable.kernel_idx, + constant_args_idx=constant_args_idx, + grid=grids, # type: ignore[arg-type] + # TMA descriptor capturing not yet + # supported in non-dynamo tracing + tma_descriptor_metadata={}, + kwargs=graphable_args, + ) + + +tracing_triton_hopifier_singleton = TracingTritonHOPifier() + + +class TraceableTritonKernelWrapper: + kernel: "TritonKernelType" + kernel_idx: Optional[int] + grid: Optional["TritonGridType"] + + def __init__( + self, + kernel: "TritonKernelType", + kernel_idx: Optional[int], + grid: Optional["TritonGridType"], + ) -> None: + self.kernel = None + self.grid = None + tracing_triton_hopifier_singleton.init_variable(self, kernel, kernel_idx, grid) + assert self.kernel is not None + + def __getitem__(self, *args: Sequence[Any]) -> "TraceableTritonKernelWrapper": + return tracing_triton_hopifier_singleton.call_getitem(self, args) # type: ignore[return-value] + + def run(self, *args: Sequence[Any], **kwargs: Dict[str, Any]) -> Any: + from torch._library.triton import is_wrap_triton_enabled + + if is_wrap_triton_enabled(): + return tracing_triton_hopifier_singleton.call_run(self, args, kwargs, None) + else: + assert self.kernel is not None + return self.kernel.run(*args, **kwargs) + + def __call__(self, *args: Sequence[Any], **kwargs: Dict[str, Any]) -> Any: + from torch._library.triton import is_wrap_triton_enabled + + if is_wrap_triton_enabled(): + return tracing_triton_hopifier_singleton.call_triton_kernel( + self, args, kwargs, None + ) + else: + assert self.kernel is not None + return self.kernel[self.grid](*args, **kwargs) + + def specialize_symbolic(self, arg: Sequence[Any]) -> Any: + import torch + + # See [Note: Specialize tl.constexpr args in user-defined triton kernels] + if isinstance(arg, (torch.SymInt, torch.SymBool, torch.SymFloat)): + return guard_scalar(arg) + return arg diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/utils.py b/lib/python3.12/site-packages/torch/_higher_order_ops/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2d00e8ec3c18e52f9c6a689e5d6c5174d978d7ca --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/utils.py @@ -0,0 +1,519 @@ +# mypy: allow-untyped-defs +import functools +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable, List, Tuple, Union + +import torch +import torch.fx.traceback as fx_traceback +import torch.utils._pytree as pytree +from torch._ops import OperatorBase +from torch.fx.experimental.proxy_tensor import make_fx +from torch.fx.passes.shape_prop import TensorMetadata +from torch.multiprocessing.reductions import StorageWeakRef + + +@dataclass +class UnsupportedAliasMutationException(RuntimeError): + reason: str + + +def autograd_not_implemented_inner( + operator: OperatorBase, delayed_error: bool, *args: Any, **kwargs: Any +) -> Any: + """If autograd is enabled and any of the arguments require grad this will either + raise an error or return a DelayedError depending on the value of delayed. + + Args: + operator: The Operator to call with the *args and **kwargs with + op_name: The name of the Operator + delayed_error: If True, return a DelayedError instead of raising an error + args: The flattened operands to the Operator + kwargs: The keyword arguments to the Operator + + Raises: + RuntimeError: If autograd is enabled and any of the arguments to the Operator + """ + with torch._C._AutoDispatchBelowAutograd(): + result = operator(*args, **kwargs) + flat_operands = pytree.arg_tree_leaves(*args) + if torch.is_grad_enabled() and any( + f.requires_grad for f in flat_operands if isinstance(f, torch.Tensor) + ): + if delayed_error: + err_fn = torch._C._functions.DelayedError( + f"Autograd not implemented for {str(operator)}", + 1, + ) + + def fake_requires_grad(tensor): + if torch.is_floating_point(tensor) or torch.is_complex(tensor): + tensor = tensor.detach() + tensor.requires_grad = True + return tensor + + return pytree.tree_map_only( + torch.Tensor, lambda x: err_fn(fake_requires_grad(x)), result + ) + else: + raise RuntimeError(f"Autograd not implemented for {str(operator)}") + return result + + +def autograd_not_implemented(op: OperatorBase, deferred_error: bool) -> Callable: + def inner(*args, **kwargs): + return autograd_not_implemented_inner(op, deferred_error, *args, **kwargs) + + return inner + + +def _maybe_run_with_interpreter(fn): + maybe_interpreted_fn = fn + if isinstance(fn, torch.fx.GraphModule) and fx_traceback.has_preserved_node_meta(): + # Running graph with interpreter is needed for propagating the stack_trace + def graph_with_interpreter(*args): + with fx_traceback.preserve_node_meta(): + return torch.fx.Interpreter(fn).run(*args) + + maybe_interpreted_fn = graph_with_interpreter + return maybe_interpreted_fn + + +def reenter_make_fx(fn): + from torch.fx.experimental.proxy_tensor import _CURRENT_MAKE_FX_TRACER + + @functools.wraps(fn) + def wrapped(*args): + assert ( + _CURRENT_MAKE_FX_TRACER is not None + ), "Cannot reenter make_fx when we're not under a make_fx tracing session" + return _CURRENT_MAKE_FX_TRACER.trace_subgraph( + _maybe_run_with_interpreter(fn), *args + ) + + return wrapped + + +def _maybe_reenter_make_fx(fn): + from torch.fx.experimental.proxy_tensor import _CURRENT_MAKE_FX_TRACER + + if _CURRENT_MAKE_FX_TRACER is not None: + return reenter_make_fx(fn) + else: + + def _maybe_make_fx_with_fake_mode(fn): + @functools.wraps(fn) + def wrapped(*args): + from torch._guards import detect_fake_mode + + fake_mode = detect_fake_mode(args) + if fake_mode is None: + # we creaeta a fake_mode here to make sure we could + # trace the graph with data-dependent calls e.g. .item() + return make_fx(fn, tracing_mode="fake")(*args) + # Tracing with real if all inputs have been fakfied + return make_fx(fn)(*args) + + return wrapped + + return _maybe_make_fx_with_fake_mode(fn) + + +@contextmanager +def _set_compilation_env(): + _old_is_tracing = torch.fx._symbolic_trace._is_fx_tracing_flag + try: + # We need to turn off the is_fx_tracing_flag. Remove this flag check from dyanmo + # once we are confident fx tracing works with dynamo. + torch.fx._symbolic_trace._is_fx_tracing_flag = False + yield + finally: + torch.fx._symbolic_trace._is_fx_tracing_flag = _old_is_tracing + + +def _detect_input_mutation(gm): + input_nodes = set() + for node in gm.graph.nodes: + if node.op == "placeholder": + input_nodes.add(node) + if node.op == "call_function": + target = node.target + if isinstance(target, torch._ops.OpOverload) and target._schema.is_mutable: + for arg in node.args: + if arg in input_nodes: + return True + + for _, module in gm.named_children(): + if isinstance(module, torch.fx.GraphModule): + if _detect_input_mutation(module): + return True + + return False + + +def _detect_input_alias(gm): + input_storages = set() + for node in gm.graph.nodes: + # We need to check existence of "val" because we reuse the logic here + # for map operator, where num_mapped_args is a scalar + # and doesn't have a "val" meta. + if ( + node.op == "placeholder" + and "val" in node.meta + and isinstance(node.meta["val"], torch.Tensor) + ): + input_storages.add(StorageWeakRef(node.meta["val"]._typed_storage())) + if node.op == "output": + + def check_alias(out): + if ( + out is not None + and "val" in out.meta + and isinstance(out.meta["val"], torch.Tensor) + ): + out_storage = StorageWeakRef(out.meta["val"]._typed_storage()) + return out_storage in input_storages + return False + + if any(pytree.tree_leaves(pytree.tree_map(check_alias, node.args))): + return True + + for _, module in gm.named_children(): + if isinstance(module, torch.fx.GraphModule) and _detect_input_alias(module): + return True + + return False + + +def has_potential_input_alias_or_mutation(gm, inputs, pre_dispatch=False): + try: + gm = make_fx(gm, pre_dispatch=pre_dispatch)(*inputs) + except UnsupportedAliasMutationException: + # this can happen when nested cond_op is + # functionalized + return True + except Exception as e: + raise e + + return _detect_input_mutation(gm) or _detect_input_alias(gm) + + +def _has_potential_branch_input_mutation(branch, inputs, pre_dispatch=False): + """ + Dispatch-trace the branch with inputs and check if + producing graph has mutable op on the input. This is + bit restrictive as the branch must be traceable. + """ + try: + gm = make_fx(branch, pre_dispatch=pre_dispatch)(*inputs) + except UnsupportedAliasMutationException: + # this can happen when nested cond_op is + # functionalized + return True + except Exception as e: + raise e + + return _detect_input_mutation(gm) + + +def _has_potential_branch_input_alias(branch, inputs, pre_dispatch=False): + """ + Dispatch-trace the branch with inputs and check if + producing graph has output aliasing the branch input. This is + bit restrictive as the branch must be traceable. + """ + try: + gm = make_fx(branch, pre_dispatch=pre_dispatch)(*inputs) + except UnsupportedAliasMutationException: + # this can happen when nested cond_op is + # functionalized + return True + except Exception as e: + raise e + + return _detect_input_alias(gm) + + +def unique_graph_id(proxy_mode, prefix): + """Returns a unique name and id for a graph to be added to a proxy_mode tracer""" + # There are probably better ways - I know that create_arg has some self incrementing name + # magic to it, but since we explicitly have to get the name for register_module, + # I was not sure how to do that. This kinda simulates it. + next_name = None + i = 0 + while not next_name: + candidate = f"{prefix}_{i}" + if hasattr(proxy_mode.tracer.root, candidate): + i += 1 + else: + next_name = candidate + return i, next_name + + +def _from_fun(t): + from torch._functorch.aot_autograd import from_fun + from torch._subclasses.functional_tensor import FunctionalTensor + + if isinstance(t, torch.Tensor): + if t.dtype != torch.bool: + return torch.empty_strided( + t.size(), + t.stride(), + dtype=t.dtype, + requires_grad=t.requires_grad, + device=t.device, + ) + else: + # clone of a functional tensor produces a functional tensor + # but we want to avoid it so we clone a non-functional version + maybe_unfunc_t = t + if isinstance(t, FunctionalTensor): + torch._sync(t) + maybe_unfunc_t = from_fun(t) + elif torch._is_functional_tensor(t): + # need to handle both types of functionalization here: + # these are the tensors that came from the user, + # which could be either FunctionalTensorWrapper or FunctionalTensor + torch._sync(t) + maybe_unfunc_t = torch._from_functional_tensor(t) + return maybe_unfunc_t.clone() + return t + + +def clone_outputs_aliasing_inputs(args): + input_storage = { + StorageWeakRef(arg._typed_storage()) + for arg in args + if isinstance(arg, torch.Tensor) + } + + def maybe_clone(t): + if ( + isinstance(t, torch.Tensor) + and StorageWeakRef(t._typed_storage()) in input_storage + ): + return t.clone() + return t + + return maybe_clone + + +def prepare_fw_with_masks(fn): + def fw_with_masks(*args): + fw_out = fn(*args) + return fw_out, [ + True if isinstance(ret, torch.Tensor) and ret.requires_grad else False + for ret in fw_out + ] + + return fw_with_masks + + +# TODO: The parameter use_output_and_grad_bw is required because some operations +# that utilize this function, such as the while_loop, may require (grad, fwd_outputs) +def create_fw_bw_graph(fn, use_output_and_grad_bw, fw_inputs, fw_outputs): + from torch._functorch.aot_autograd import AOTConfig, create_joint + + # Note:[HOP create fw_bw graph] We create "clean" environments for make_fx by suspending all dispatch keys + # between Autograd and Python key. Currently, we only suspend functionalization but more can be + # added when required. Will encounter two problems if we don't suspend functionalization: + # + # 1. make_fx fails to capture operations on input: the inputs are wrapped as _to_functional_tensor_wrapper, + # but they will be unwrapped before entering ProxyTorchDispatchMode as part of the dispatching. + # However, it's the outside wrapper that tracer creates proxies for. This casuses tracer fail to + # fetch the proxy for the inputs and fail to capture any operations on them. + # + # 2. make_fx fails to capture output: the outputs after ProxyTorchDispatchMode are further + # wrapped as FunctionalTensorWrapper in Functionalize key after return. However, the tracer + # only associates the inner tensor with proxy in ProxyTorchDispatchMode. Therefore, + # when creating the output node, it fails to associate the wrapped tensor with its proxy. + # Instead, it will create _tensor_constant as output. + + dummy_aot_config = AOTConfig( + fw_compiler=None, # type: ignore[arg-type] + bw_compiler=None, # type: ignore[arg-type] + partition_fn=None, # type: ignore[arg-type] + decompositions={}, + num_params_buffers=0, + aot_id=0, + keep_inference_input_mutations=False, + ) + + example_grad = [_from_fun(out) for out in fw_outputs] + num_grads = len(example_grad) + fw_graph = _maybe_reenter_make_fx(fn)(*fw_inputs) + + def joint_fn(*joint_operands_grads): + if use_output_and_grad_bw: + grads = joint_operands_grads[0] + inputs = joint_operands_grads[1][-1:] + else: + grads = joint_operands_grads[:num_grads] + inputs = joint_operands_grads[num_grads:] + + joint = create_joint(prepare_fw_with_masks(fn), aot_config=dummy_aot_config) + _, grads = joint( + list(inputs), + [grad for grad in grads if grad is not None and grad.requires_grad], + ) + + # In order to keep map functional for backward graph, + # we clone outputs that are aliasing inputs + maybe_clone = clone_outputs_aliasing_inputs(joint_operands_grads) + + return pytree.tree_map(maybe_clone, grads) + + if use_output_and_grad_bw: + example_xs_out = list(fw_inputs) + list(fw_outputs) + joint_graph = _maybe_reenter_make_fx(joint_fn)( + (list(example_grad), list(example_xs_out)) + ) + else: + example_xs_out = list(fw_inputs) + joint_graph = _maybe_reenter_make_fx(joint_fn)( + *(list(example_grad) + list(example_xs_out)) + ) + + return fw_graph, joint_graph + + +def _unstack_pytree(xs): + flat_xs, inspec = pytree.tree_flatten(xs) + if not all(isinstance(xs, torch.Tensor) for xs in flat_xs): + raise RuntimeError(f"Leaves of xs must be Tensor {flat_xs}") + + if not all(xs.shape[0] == flat_xs[0].shape[0] for xs in flat_xs): + raise RuntimeError( + f"Leaves of xs must have same leading dimension size {[xs.shape for xs in flat_xs]}" + ) + + a = zip(*flat_xs) + + pytrees = [pytree.tree_unflatten(tuple, inspec) for tuple in a] + return pytrees + + +def _stack_pytree(pytrees): + flat_out = [] + out_spec = None + for pt in pytrees: + flat_pt, out_spec = pytree.tree_flatten(pt) + flat_out.append(flat_pt) + assert out_spec is not None + b = zip(*flat_out) + stacked_out = [] + for leaves in b: + if all(isinstance(leaf, torch.Tensor) for leaf in leaves): + stacked_out.append(torch.stack(leaves)) + elif all(leaf is None for leaf in leaves): + # Backward graph can return None output when forward inputs doesn't require grad. + # When we eagerly execute backward graph, we need to call _stack_pytree on its output, + # therefore we need to deal with None output. + stacked_out.append(None) # type: ignore[arg-type] + else: + raise RuntimeError(f"Cannot stack {leaves}.") + return pytree.tree_unflatten(stacked_out, out_spec) + + +# We cannot call save_for_backward for symints. This helper function +# can be used to save symints as direct attributes of ctx in autograd.Function. +# +# For example, if args = (x, y, s0, z, s1), +# save_tensors_and_symints_for_backward will partition the args into two lists, and a bookkeeping list pos: +# partitioned_args[0] = (x, y, z) +# partitioned_args[1] = (s0, s1) +# pos = (0, 0, 1, 0, 1) +# pos list keeps track of which partition the args +# is partitioned into in order to recover it in saved_tensors_and_symints. +# +# In saved_tensors_and_symints, we can recover the original args by: +# iterating over the pos list and pop one item from the front of paritioned_args[pos[i]]. +# We use t_idx and s_idx to keep track of the next index of the item we are going to pop for the two lists. +def save_tensors_and_symints_for_backward(ctx, args): + assert all( + isinstance(arg, (torch.Tensor, torch.SymInt, int, type(None))) for arg in args + ), args + partitioned_args: List[Any] = [[], []] + pos = [] + for i, arg in enumerate(args): + idx = 0 if isinstance(arg, torch.Tensor) else 1 + partitioned_args[idx].append(arg) + pos.append(idx) + + assert not hasattr(ctx, "sym_int_args"), "ctx already has sym_int_args attribute." + assert not hasattr(ctx, "pos"), "ctx already has pos attribute." + ctx.save_for_backward(*partitioned_args[0]) + ctx.sym_int_args = partitioned_args[1] + ctx.pos = pos + + +def saved_tensors_and_symints(ctx): + args = [] + t_idx = 0 + s_idx = 0 + saved_tensors = ctx.saved_tensors + for p in ctx.pos: + if p == 0: + args.append(saved_tensors[t_idx]) + t_idx += 1 + else: + args.append(ctx.sym_int_args[s_idx]) + s_idx += 1 + assert t_idx + s_idx == len(ctx.pos) + return tuple(args) + + +def get_dummy_aot_autograd_config(): + from torch._functorch.aot_autograd import AOTConfig + + return AOTConfig( + fw_compiler=None, # type: ignore[arg-type] + bw_compiler=None, # type: ignore[arg-type] + partition_fn=None, # type: ignore[arg-type] + decompositions={}, + num_params_buffers=0, + aot_id=0, + keep_inference_input_mutations=False, + ) + + +# Slices off the first element of a given dimension +def first_slice_copy(t: torch.Tensor, dim: int = 0) -> torch.Tensor: + return torch.select_copy(t, dim, 0) + + +# Reports the difference between meta of two tensors in a string +def diff_tensor_meta( + meta1: TensorMetadata, meta2: TensorMetadata, check_grad=True +) -> List[str]: + from torch.fx.experimental.symbolic_shapes import GuardOnDataDependentSymNode + + pair_diffs = [] + for meta_name in TensorMetadata._fields: + if not check_grad and meta_name == "requires_grad": + continue + val1 = getattr(meta1, meta_name) + val2 = getattr(meta2, meta_name) + try: + if val1 != val2: + pair_diffs.append(f"'{meta_name}: {val1} vs {val2}'") + except GuardOnDataDependentSymNode as _: + pair_diffs.append(f"'{meta_name}: {val1} vs {val2}'") + continue + return pair_diffs + + +# Note [lifted arg types in hop] +# For dynamoed hops, we automatically lift the free symbols in tensors as arguments. +# This has implications for the types of lifted args for different dispatch keys: +# 1. functionalization, FakeTensorMode, ProxyTorchDispatchMode, Autograd need to support torch.Symint +# lifted args because it's on the path of torch.compile(dynamic=True). +# 2. functionalization, FakeTensorMode, ProxyTorchDispatchMode, Autograd, CompositeExplicitAutograd need +# to support int arguments. In the eager run case, we re-trace the subgraph in AutogradKey, so inner +# hops may receive int inputs from the shape of outer tensor inputs. +# However, CompositeExplicitAutograd won't receive SymInt inputs because it only accepts real tensor inputs. +def validate_subgraph_args_types(lifted_args: Union[Tuple[Any, ...], List[Any]]): + allowed_types = (torch.Tensor, int, torch.SymInt) + assert all( + isinstance(arg, (torch.Tensor, int, torch.SymInt)) for arg in lifted_args + ), f"{lifted_args} can only be of {allowed_types} but got {tuple(type(arg) for arg in lifted_args)}" diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/while_loop.py b/lib/python3.12/site-packages/torch/_higher_order_ops/while_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..c0811ae4629f084b4a1c7325c7f0b88bfd53f9f5 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/while_loop.py @@ -0,0 +1,340 @@ +# mypy: allow-untyped-defs +from typing import Callable, List, Tuple, Union + +import torch +import torch.utils._pytree as pytree +from torch._C import DispatchKey +from torch._higher_order_ops.utils import ( + _has_potential_branch_input_alias, + _has_potential_branch_input_mutation, + _maybe_run_with_interpreter, + _set_compilation_env, + autograd_not_implemented, + diff_tensor_meta, + reenter_make_fx, + UnsupportedAliasMutationException, + validate_subgraph_args_types, +) +from torch._ops import HigherOrderOperator +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + _temp_remove_metadata_torch_function_mode, + ProxyTorchDispatchMode, + track_tensor_tree, +) +from torch.fx.passes.shape_prop import _extract_tensor_metadata + + +class WhileLoopOp(HigherOrderOperator): + def __init__(self) -> None: + super().__init__("while_loop") + + def __call__( + self, + cond_fn: Callable, + body_fn: Callable, + carried_inputs: Tuple[Union[torch.Tensor, int, float, bool]], + additional_inputs: Tuple[Union[torch.Tensor, torch.SymInt, int], ...], + /, + ): + if not isinstance(carried_inputs, tuple): + raise RuntimeError( + f"carried_inputs must be a tuple, got {type(carried_inputs)}" + ) + if not isinstance(additional_inputs, tuple): + raise RuntimeError( + f"additional_inputs must be a tuple, got {type(additional_inputs)}" + ) + + validate_subgraph_args_types(carried_inputs) + validate_subgraph_args_types(additional_inputs) + return super().__call__(cond_fn, body_fn, carried_inputs, additional_inputs) + + +while_loop_op = WhileLoopOp() + + +def while_loop(cond_fn, body_fn, carried_inputs): + r""" + Run body_fn(*carried_inputs) while cond_fn(*carried_inputs) returns a True scalar tensor. Returns the output of body_fn or + initial carried_inputs. + + .. warning:: + `torch.while_loop` is a prototype feature in PyTorch. It has limited support for input and output types and + doesn't support training currently. Please look forward to a more stable implementation in a future version of PyTorch. + Read more about feature classification at: https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype + + `while_loop` is a structured control flow operator. It preserves the loop semantic across the torch.compile and torch.export. + + `while_loop` is equivalent to the following: + + def while_loop(cond_fn, body_fn, carried_inputs): + val = carried_inputs + while cond_fn(*val): + val = body_fn(*val) + return val + + Args: + cond_fn (Callable): A callable function that returns a boolean Scalar tensor. + + body_fn (Callable): A callable function that takes the same inputs as `cond_fn` and returns a tuple of tensors + + carried_inputs (Tuple of possibly nested dict/list/tuple of tensors): A tuple of inputs to cond_fn and body_fn. It's also + the initial value of states that are carried across iterations. + + Example: + + def cond_fn(iter, x): + return iter.sum() < 10 + + def body_fn(iter, x): + return iter + 1, x.sin() + + while_loop(cond_fn, body_fn, (torch.zeros(1), torch.randn(3, 4))) + + Restrictions: + + - body_fn must return tensors with the same metadata (e.g.shape, dtype) as inputs. + + - body_fn and cond_fn must not in-place mutate the carried_inputs. A clone before the mutation is required. + + - body_fn and cond_fn must not mutate python varialbles (e.g. list/dict) created outside of the body_fn. + + - body_fn and cond_fn's output cannot aliase any of the inputs. A clone is required. + + .. warning:: + Temporal Limitations: + + - 'while_loop' only supports **inference** right now. Autograd will be supported in the future. + + """ + from torch._dynamo.backends.debugging import ( + make_eager_backend_with_torch_function_mode, + ) + + # Currently, additional_inputs is not a user-facing input. It will be automatically set in dynamo. + # parameters and buffers accessed in cond_fn or body_fn or tensor closures will become additional_inputs. + additional_inputs: Tuple = () + + # The reason we flatten the output before calling into dynamo is that + # we want to create a consistent input ordering for cond_fn and body_fn. + # and we also want to the input ordering matches the output ordering. + # Also see NOTE: [why we cannot use "automatic" for while_loop] + # Construct flat cond_fn and flat_body_fn, which takes flattened inputs + flat_inputs, in_spec = pytree.tree_flatten((carried_inputs, additional_inputs)) + + def flat_cond_fn(*flat_args): + carried, additional = pytree.tree_unflatten(flat_args, in_spec) + return cond_fn(*carried, *additional) + + def flat_body_fn(*flat_args): + carried, additional = pytree.tree_unflatten(flat_args, in_spec) + return body_fn(*carried, *additional) + + if torch.compiler.is_dynamo_compiling(): + return while_loop_op(flat_cond_fn, flat_body_fn, tuple(flat_inputs), tuple()) + + def _validate_input(cond_fn, body_fn, carried_inputs): + from torch._higher_order_ops.utils import validate_subgraph_args_types + + if not callable(cond_fn) or not callable(body_fn): + raise RuntimeError("Expect cond_fn and body_fn to be callable.") + + validate_subgraph_args_types(flat_inputs) + + if not pytree.tree_all(lambda t: isinstance(t, torch.Tensor), carried_inputs): + raise RuntimeError( + "Expect carried_inputs to be a tuple of possibly nested dict/list/tuple that only" + f"consists of tensor leaves, but got {carried_inputs}." + ) + + _validate_input(cond_fn, body_fn, carried_inputs) + + # Dynamo is expecting a callable with "__code__" attribute. + # We cannot directly pass cond_op to it. So we wrap it in a dummy function. + def _while_loop_op_wrapper(*args, **kwargs): + return while_loop_op(*args, **kwargs) + + with _set_compilation_env(), torch._dynamo.utils.disable_cache_limit(): + with _temp_remove_metadata_torch_function_mode() as metadata_mode: + with _temp_remove_metadata_torch_function_mode() as metadata_mode: + if metadata_mode: + backend = make_eager_backend_with_torch_function_mode(metadata_mode) + else: + backend = "eager" + return torch.compile( + _while_loop_op_wrapper, backend=backend, fullgraph=True + )(flat_cond_fn, flat_body_fn, tuple(flat_inputs), tuple()) + + +@while_loop_op.py_impl(DispatchKey.CompositeExplicitAutograd) +def while_loop_dense(cond_fn, body_fn, carried_inputs, additional_inputs): + carried_vals = carried_inputs + + def _is_boolean_scalar_tensor(pred): + return ( + isinstance(pred, torch.Tensor) + and pred.size() == torch.Size([]) + and pred.dtype == torch.bool + ) + + if not isinstance(carried_inputs, tuple): + raise RuntimeError( + f"carried_inputs must be a tuple but got {type(carried_inputs)}" + ) + + while pred := cond_fn(*carried_vals, *additional_inputs): + if not _is_boolean_scalar_tensor(pred): + raise RuntimeError( + f"cond_fn must return a boolean scalar tensor but got {pred}" + ) + out = body_fn(*carried_vals, *additional_inputs) + assert isinstance( + out, tuple + ), f"body_fn should return a tuple but got {type(out)}" + assert len(out) == len( + carried_inputs + ), "body_fn should return the same number of elements as carried_inputs" + carried_vals = out + return carried_vals + + +while_loop_op.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(while_loop_op, deferred_error=True) +) + + +@while_loop_op.py_impl(ProxyTorchDispatchMode) +def while_loop_tracing(mode, cond_fn, body_fn, carried_inputs, additional_inputs): + def _trace_while_loop( + proxy_mode, while_loop_op, cond_fn, body_fn, carried_inputs, additional_inputs + ): + cond_graph = reenter_make_fx(cond_fn)(*carried_inputs, *additional_inputs) + body_graph = reenter_make_fx(body_fn)(*carried_inputs, *additional_inputs) + + next_name = None + i = 0 + while not next_name: + candidate = f"while_loop_cond_graph_{i}" + if hasattr(proxy_mode.tracer.root, candidate): + i += 1 + else: + next_name = candidate + cond_graph_name = next_name + body_graph_name = f"while_loop_body_graph_{i}" + assert not hasattr(proxy_mode.tracer.root, body_graph_name) + + proxy_mode.tracer.root.register_module(cond_graph_name, cond_graph) + proxy_mode.tracer.root.register_module(body_graph_name, body_graph) + + args = (cond_graph, body_graph, carried_inputs, additional_inputs) + + proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, args) + + out_proxy = proxy_mode.tracer.create_proxy( + "call_function", while_loop_op, proxy_args, {}, name="while_loop" + ) + + out = while_loop_op(cond_graph, body_graph, carried_inputs, additional_inputs) + return track_tensor_tree( + out, out_proxy, constant=None, tracer=proxy_mode.tracer + ) + + return _trace_while_loop( + mode, while_loop_op, cond_fn, body_fn, carried_inputs, additional_inputs + ) + + +def check_outputs_carry_consistency( + outs: List[torch.Tensor], carries: List[torch.Tensor] +) -> None: + all_diffs_in_meta = [] + for out, cry in zip(outs, carries): + if diff := diff_tensor_meta( + _extract_tensor_metadata(cry), _extract_tensor_metadata(out) + ): + all_diffs_in_meta.append(",".join(diff)) + if all_diffs_in_meta: + diff_str = "\n".join(all_diffs_in_meta) + raise RuntimeError( + f"Expected carried_inputs and body outputs return tensors with same metadata but found:\n{diff_str}" + ) + + +@while_loop_op.py_impl(FakeTensorMode) +def while_loop_fake_tensor_mode( + mode, cond_fn, body_fn, carried_inputs, additional_inputs +): + with mode: + # NOTE: [Handling unback symints created in subgraph of while_loop] + # The idea is that the scope of unbacked symints are limited to the subgraph. + # + # We're implementing the fake tensor mode of while_loop operator. + # and we run body_fn once to get an fake output. + # Let's only consider tensor output for now: + # + # Case 1: + # if the unbacked symints is local to the subgraph e.g. + # def body_fn(it, x): + # nz = x.nonzero() + # return it+1. nz.sum() + # we can just ignore the newly created unbacked symints because it has + # no effect on the output of while_loop and it's tracked when we tracing. + # the subgraph. + # + # Case 2.1: + # if the unbacked symints are part of output of while_loop e.g. + # def body_fn(it, x): + # nz = x.nonzero() + # return it+1, nz + # This will fail the shape check because in each iteration, the carried_input's shape + # must match the output shape as nz.shape contains newly allocated unbacked symint, this + # won't match the carried_input's shape. + # + # Case 2.2: + # if the unbacked symints are part of carried_inputs e.g. + # nz = a.nonzero() + # body_fn(it, nz): + # return it+1. nz.sin() + 1, + # There's no new unbacked symints allocated in subgraph, so we're safe. + with mode.shape_env.ignore_fresh_unbacked_symbols(): + # body_fn return output with the same pytree and tensor meta data as carried_inputs + # so we could just return the output after one iteration. + body_outs = body_fn(*carried_inputs, *additional_inputs) + check_outputs_carry_consistency(body_outs, carried_inputs) + return body_outs + + +@while_loop_op.py_functionalize_impl +def while_loop_func(ctx, cond_fn, body_fn, carried_inputs, additional_inputs): + unwrapped_carried_inputs = ctx.unwrap_tensors(carried_inputs) + unwrapped_additional_inputs = ctx.unwrap_tensors(additional_inputs) + unwrapped_inputs = unwrapped_carried_inputs + unwrapped_additional_inputs + with ctx.redispatch_to_next() as m: + functional_cond_fn = ctx.functionalize(_maybe_run_with_interpreter(cond_fn)) + functional_body_fn = ctx.functionalize(_maybe_run_with_interpreter(body_fn)) + pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch + for fn, fn_name in [ + (functional_cond_fn, "cond_fn"), + (functional_body_fn, "body_fn"), + ]: + if _has_potential_branch_input_mutation( + fn, unwrapped_inputs, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException( + f"torch.while_loop's {fn_name} might be modifying the input!" + ) + + if _has_potential_branch_input_alias( + fn, unwrapped_inputs, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException( + f"torch.while_loop's {fn_name} might be aliasing the input!" + ) + ret = while_loop_op( + functional_cond_fn, + functional_body_fn, + unwrapped_carried_inputs, + unwrapped_additional_inputs, + ) + return ctx.wrap_tensors(ret) diff --git a/lib/python3.12/site-packages/torch/_higher_order_ops/wrap.py b/lib/python3.12/site-packages/torch/_higher_order_ops/wrap.py new file mode 100644 index 0000000000000000000000000000000000000000..9310c55ddacffa398d13046532a7644b265c1c3f --- /dev/null +++ b/lib/python3.12/site-packages/torch/_higher_order_ops/wrap.py @@ -0,0 +1,240 @@ +# mypy: allow-untyped-defs +import inspect +import itertools +import logging +from typing import Optional + +from torch._logging import warning_once +from torch._ops import HigherOrderOperator +from torch.types import _dtype +from torch.utils.checkpoint import checkpoint, CheckpointPolicy + + +log = logging.getLogger(__name__) + +uid = itertools.count(1) + + +# Used for testing the HigherOrderOperator mechanism +class Wrap(HigherOrderOperator): + def __init__(self) -> None: + super().__init__("wrap") + + def __call__(self, func, *args, **kwargs): + # Dynamo already traces the body of HigherOrderOp beforehand when it + # so no need to trace into it. + import torch._dynamo # noqa: F401 + from torch._dynamo import disable + + @disable + def wrapper(): + result = func(*args, **kwargs) + return result + + return wrapper() + + +wrap = Wrap() + + +class WrapWithSetGradEnabled(HigherOrderOperator): + def __init__(self) -> None: + super().__init__("wrap_with_set_grad_enabled") + + def __call__(self, enable_grad, wrapped_func, *args, **kwargs): + # Dynamo already traces the body of HigherOrderOp beforehand when it + # so no need to trace into it. + import torch._dynamo # noqa: F401 + from torch._dynamo import disable + + @disable + def wrapper(): + with torch.set_grad_enabled(enable_grad): + return wrapped_func(*args, **kwargs) + + return wrapper() + + +wrap_with_set_grad_enabled = WrapWithSetGradEnabled() + + +class WrapWithAutocast(HigherOrderOperator): + def __init__(self): + super().__init__("wrap_with_autocast") + + def __call__( + self, + device_type: str, + dtype: Optional[_dtype], + enabled: bool, + cache_enabled: Optional[bool], + wrapped_func, + *args, + **kwargs, + ): + # Dynamo already traces the body of HigherOrderOp beforehand when it + # so no need to trace into it. + import torch._dynamo # noqa: F401 + from torch._dynamo import disable + + @disable + def wrapper(): + with torch.autocast(device_type, dtype, enabled, cache_enabled): + return wrapped_func(*args, **kwargs) + + return wrapper() + + +wrap_with_autocast = WrapWithAutocast() + + +class WrapActivationCheckpoint(HigherOrderOperator): + """ + This operator is used to wrap torch.utils.checkpoint. This avoids + TorchDynamo to look into saved tensor hooks and directly passes the control + to AOT Autograd, which is ok with tracing saved tensor hooks. As a result of + AOT tracing torch.utils.checkpoint code, we have a backward graph with + recomputed forward nodes. + + However, we might deprecate this operator soon. The difficulty arises in the + functionalization of rng ops. Today, there are two different + functionalization of rng ops - one at AOT autograd and other at Inductor. + And they are difficult to map to each other. The rng states also complicate + pattern matching in Inductor. Due to the ease of implementation, we are + currently inclined towards functionalization at Inductor level, which means + that duplication/recomputation is done as a compiler pass in the + partitioners. See TagActivationCheckpoint for more information. + """ + + def __init__(self) -> None: + super().__init__("wrap_activation_checkpoint", cacheable=False) + + def __call__(self, function, *args, **kwargs): + # use_reentrant is set to False because this op is going to be traced. + # And we ensure that AOT Autograd traces through the non reentrant + # version of checkpointing. + import torch.fx.traceback as fx_traceback + from torch.fx import Interpreter + + kwargs["use_reentrant"] = False + kwargs["preserve_rng_state"] = False + # Using interpreter allows preservation of metadata through torch.compile stack. + with fx_traceback.preserve_node_meta(): + return checkpoint(Interpreter(function).run, *args, **kwargs) + + +wrap_activation_checkpoint = WrapActivationCheckpoint() + + +class TagActivationCheckpoint(HigherOrderOperator): + """ + This operator is supposed to be used only with torch.compile stack. This + accepts a Fx graph module which needs to be checkpointed. This operator adds + "recomputable" tag to the nodes of the Fx graph that should be recomputed. + + The goal is to: + 1. Avoid using Dynamo to trace through saved tensor hooks. + 2. For selective checkpointing case, let AOTAutograd trace through + saved tensor hooks but has special logic with TorchDispatchMode to override + the usual saved_tensor_hooks fn logic in order to tag the nodes. + 3. Rely on the partitioners to actually duplicate the nodes. + This sits well in the torch.compile stack, because by the time graph + reaches partitioner, inductor has already run its functionalization of rng + ops (by setting fixed seed for each random op, see `replace_random_passes`). + Therefore, the duplication of nodes, by design, respects the rng states in + the forward and recomputed forward in backward. + """ + + def __init__(self) -> None: + super().__init__("tag_activation_checkpoint", cacheable=False) + + @staticmethod + def divide_kwargs(kwargs): + """ + checkpoint fn can have mixed kwargs between checkpointed fn and + checkpoint fn itself. For example + >> def gn(x, y, z=None): + >> a = torch.matmul(x, y) + >> if z is not None: + >> return torch.matmul(a, z) + >> return a + >> def fn(x, y, z): + >> return torch.cos(checkpoint(gn, x, y, use_reentrant=False, z=z)) + In the above case, z belongs to checkpointed function gn, but + use_reentrant belongs to the checkpoint function. This function splits + the kwargs into checkpoint_kwargs and gmod_kwargs (or + checkpointed_fn_kwargs). + We do sorting to ensure same graph from run to run for better + debuggability. It is not required for correctness. + """ + ckpt_signature = inspect.signature(checkpoint) + checkpoint_keys = set() + for name in ckpt_signature.parameters: + if name in ("function", "args", "kwargs"): + continue + checkpoint_keys.add(name) + + # `preserve_rng_state` is not a regular kwarg + checkpoint_keys.add("preserve_rng_state") + + checkpoint_kwargs = { + name: kwargs[name] for name in kwargs.keys() if name in checkpoint_keys + } + gmod_kwargs = { + name: kwargs[name] for name in kwargs.keys() if name not in checkpoint_keys + } + return checkpoint_kwargs, gmod_kwargs + + def tag_nodes(self, gmod, is_sac): + unique_graph_id = next(uid) + for node in gmod.graph.nodes: + if node.op in ("call_function", "call_method", "call_module"): + node.meta["ac_graph_id"] = unique_graph_id + if is_sac: + # For selective checkpointing, we will populate this tag later in _CachingTorchDispatchMode. + node.meta["recompute"] = None + else: + # Under vanilla activation checkpointing, all nodes should be recomputed. + node.meta["recompute"] = CheckpointPolicy.PREFER_RECOMPUTE + return gmod + + def __call__(self, gmod, *args, **kwargs): + import torch.fx.traceback as fx_traceback + from torch.fx import Interpreter + + if "_checkpoint_context_fn" in gmod.meta: + warning_once( + log, + """ +Detected that context_fn is passed to torch.utils.checkpoint under torch.compile. +Please make sure the checkpointed region does not contain in-place ops (e.g. torch.relu_). +""", + ) + # use_reentrant is set to False because this op is going to be traced. + # And we ensure that AOT Autograd traces through the non reentrant + # version of checkpointing. + kwargs["use_reentrant"] = False + # preserve_rng_state is set to False because we want to prevent AOTAutograd from tracing through + # `torch.random.fork_rng` op (which is not supported yet under CUDA). + # This doesn't mean that we don't preserve RNG state. Instead, we will always preserve RNG state + # regardless of this flag (by doing RNG functionalization via `replace_random_passes` in Inductor + # instead of in AOTAutograd). + kwargs["preserve_rng_state"] = False + kwargs["context_fn"] = gmod.meta["_checkpoint_context_fn"] + # We first tag all nodes as "recompute" in this graph, and then we undo the "recompute" tag + # for specific nodes in _CachingTorchDispatchMode in torch/utils/checkpoint.py. + gmod = self.tag_nodes(gmod, is_sac=True) + # Using interpreter allows preservation of metadata through torch.compile stack. + with fx_traceback.preserve_node_meta(): + return checkpoint(Interpreter(gmod).run, *args, **kwargs) + else: + gmod = self.tag_nodes(gmod, is_sac=False) + # Using interpreter allows preservation of metadata through torch.compile stack. + # TODO: We want to use the same `checkpoint(Interpreter(gmod).run, *args, **kwargs)` here + # as the `context_fn != None` case, but that depends on in-place op support in TorchDispatchMode + torch.compile. + # (for details on in-place op issue, run `test_compile_selective_checkpoint_inplace_op` unit test) + with fx_traceback.preserve_node_meta(): + return Interpreter(gmod).run(*args) + + +tag_activation_checkpoint = TagActivationCheckpoint() diff --git a/lib/python3.12/site-packages/torch/_prims/__init__.py b/lib/python3.12/site-packages/torch/_prims/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9f0434c4d6cb9a646a7200cc3249f3f695b40b72 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_prims/__init__.py @@ -0,0 +1,2946 @@ +# mypy: allow-untyped-defs +import operator +from enum import Enum +from functools import partial, reduce +from typing import Callable, List, Optional, Sequence, Tuple, Type, Union + +import torch +import torch._prims_common as utils +import torch.library +from torch import sym_float, Tensor +from torch._C import _get_default_device +from torch._higher_order_ops.effects import new_token_tensor +from torch._library.utils import is_functional_schema +from torch._prims.debug_prims import register_debug_prims +from torch._prims.rng_prims import register_rng_prims +from torch._prims_common import ( + Dim, + DimsSequenceType, + DimsType, + IntLike, + Number, + NumberType, + RETURN_TYPE, + ShapeType, + StrideType, + TensorLike, + TensorLikeType, + type_to_dtype, +) +from torch._prims_common.wrappers import backwards_not_supported +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode +from torch.overrides import handle_torch_function, has_torch_function +from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten + + +prim = torch.library.Library("prims", "DEF") +prim_impl = torch.library.Library("prims", "IMPL", "CompositeExplicitAutograd") +prim_backend_select_impl = torch.library.Library("prims", "IMPL", "BackendSelect") +prim_autograd_impl = torch.library.Library("prims", "IMPL", "Autograd") +prim_meta_impl = torch.library.Library("prims", "IMPL", "Meta") + +# Experimental module containing prototype "primitive" operations. + +__all__ = [ + # + # Common datastructures and helpers + # + "RETURN_TYPE", + # + # Elementwise unary prims + # + "abs", + "acos", + "acosh", + "asin", + "asinh", + "atan", + "atanh", + "cos", + "cosh", + "bessel_i0", + "bessel_i0e", + "bessel_i1", + "bessel_i1e", + "bessel_j0", + "bessel_j1", + "bitwise_not", + "cbrt", + "ceil", + "conj_physical", + "digamma", + "erf", + "erf_inv", + "erfc", + "erfcx", + "exp", + "expm1", + "exp2", + "fill", + "floor", + "imag", + "isfinite", + "lgamma", + "log", + "log1p", + "log2", + "log10", + "ndtri", + "neg", + "real", + "reciprocal", + "round", + "sign", + "signbit", + "sin", + "sinh", + "spherical_bessel_j0", + "sqrt", + "tan", + "tanh", + "trunc", + # + # Elementwise binary prims + # + "add", + "atan2", + "bitwise_and", + "bitwise_or", + "bitwise_xor", + # 'complex', # needs custom meta + "div", + "eq", + "fmax", + "fmin", + "fmod", + "frexp", + "gcd", + "ge", + "gt", + "hypot", + "igamma", + "igammac", + "le", + "lt", + "maximum", + "minimum", + "mul", + "ne", + "nextafter", + "pow", + "remainder", + "rsqrt", + "shift_left", + "shift_right_arithmetic", + "shift_right_logical", # not implemented + "sub", + "zeta", + # + # View prims + # + "as_strided", + "broadcast_in_dim", + "collapse_view", + "conj", + "expand_dims", + "slice", + "slice_in_dim", # implemented using slice -- make this a ref? + "split_dim", + "squeeze", + "transpose", + "view_of", + "view_element_type", + # + # Functionalized view mutations + # + "as_strided_scatter", + # + # Shape prims + # + "collapse", + "cat", + "reshape", + "rev", + # + # Conditional prims + # + "where", + # + # Data conversion and movement prims + # + "clone", + "convert_element_type", + "device_put", + "item", + "maximum_value", + "minimum_value", + "copy_strided", + # + # Inplace prims + # + "copy_to", + "resize", + # "_set", # Commented out, see note below + # + # Reduction prims + # + "amax", + "amin", + "prod", + "sum", + "xor_sum", + "var", + # + # Tensor Creation Prims + # + "empty_strided", + "empty_permuted", + "scalar_tensor", + "iota", + # + # Linear algebra (linalg) Prims + # + "svd", + # + # Randomness Prims + # + "normal", + "_uniform_helper", + # + # FFT prims + # + "fft_r2c", + "fft_c2c", + "fft_c2r", + # + # prims for making/sinking tokens + # + "_make_token", + "_sink_tokens", +] + + +def TensorMeta( + tensorlike: Optional[Union[NumberType, torch.Tensor]] = None, + *, + shape: Optional[ShapeType] = None, + strides: Optional[StrideType] = None, + dtype: Optional[torch.dtype] = None, + device: Optional[Union[torch.device, str]] = None, +): + if isinstance(tensorlike, Number): + assert not shape and (shape is None or isinstance(shape, Sequence)) + assert not strides and (strides is None or isinstance(strides, Sequence)) + inferred_shape: Tuple[int, ...] = () + inferred_strides: Tuple[int, ...] = () + inferred_dtype = type_to_dtype(type(tensorlike)) + inferred_device = torch.device("cpu") + # TODO: This looks wrong, a number that is wrapped into a tensor + # needs to behave differently than a scalar tensor for type + # promotion purposes + elif tensorlike is not None: + assert isinstance(tensorlike, torch.Tensor) + inferred_shape = tuple(tensorlike.shape) + inferred_strides = tuple(tensorlike.stride()) + inferred_dtype = tensorlike.dtype + inferred_device = tensorlike.device + else: + # If no tensorlike "example" is given then all metadata + # must be provided explicitly + assert shape is not None + assert strides is not None + assert dtype is not None + assert device is not None + + shape = inferred_shape if shape is None else tuple(shape) # type: ignore[possibly-undefined] + strides = inferred_strides if strides is None else tuple(strides) # type: ignore[possibly-undefined] + dtype = inferred_dtype if dtype is None else dtype # type: ignore[possibly-undefined] + device = inferred_device if device is None else device # type: ignore[possibly-undefined] + + if isinstance(device, str): + device = torch.device(device) + + return torch.empty_strided(shape, strides, dtype=dtype, device=device) + + +def _make_prim( + *, + schema: str, + return_type: Union[RETURN_TYPE, Tuple[RETURN_TYPE, ...]], + meta: Callable, + impl_aten: Callable, + doc: str, + tags: Optional[Sequence[torch.Tag]] = None, + use_old_custom_ops_api: bool = False, + register_conj_neg_fallthrough: bool = False, +): + """ + Creates a primitive operation. + + """ + + def _prim_impl(*args, **kwargs): + # always run the meta function because aten implementation will + # typically accept more inputs (e.g., it will do promotion and + # broadcasting) which we want to reject + meta(*args, **kwargs) + return impl_aten(*args, **kwargs) + + # Right now prims don't support autograd (we can and should add an + # argument that provides an implementation for backward here.) Because we + # don't have derivative formulas, we must setup a custom autograd function + # that raises an error if backwards is invoked + def _autograd_impl(*args, **kwargs): + return backwards_not_supported(_prim)(*args, **kwargs) + + def _backend_select_impl(*args, **kwargs): + if kwargs.get("device") and kwargs["device"].type == "meta": + return meta(*args, **kwargs) + if any(isinstance(x, torch.device) and x.type == "meta" for x in args): + return meta(*args, **kwargs) + else: + return _prim_impl(*args, **kwargs) + + name = schema.split("(")[0] + schema = schema[len(name) :] + + # register non-functional ops with old custom ops API + cpp_schema = torch._C.parse_schema(name + schema) + if use_old_custom_ops_api or not is_functional_schema(cpp_schema): + prim.define(name + schema, tags=torch.Tag.pt2_compliant_tag) + prim_impl.impl(name, _prim_impl) + prim_autograd_impl.impl(name, _autograd_impl) + prim_meta_impl.impl(name, meta) + else: + mutates_args = [ + arg.name + for arg in cpp_schema.arguments + if arg.alias_info is not None and arg.alias_info.is_write + ] + prim_def = torch.library.custom_op( + "prims::" + name, + _prim_impl, + mutates_args=tuple(mutates_args), + schema=schema, + ) + prim_def.register_fake(meta) + + # all view ops get conj/neg fallthroughs + if return_type == RETURN_TYPE.VIEW or register_conj_neg_fallthrough: + prim_def._lib.impl(name, torch.library.fallthrough_kernel, "Conjugate") + prim_def._lib.impl(name, torch.library.fallthrough_kernel, "Negative") + + _prim_packet = getattr(torch._ops.ops.prims, name) + _prim = _prim_packet.default + if tags: + _prim._tags = tags + elif aten_packet := getattr(torch.ops.aten, name, None): + overload_tags = [ + getattr(aten_packet, overload).tags for overload in aten_packet.overloads() + ] + tags_intersection = set(overload_tags[0]) + tags_intersection.intersection_update(*overload_tags[1:]) + + # dont inadvertently add to prim ops + tags_intersection.discard(torch.Tag.core) + # causes errors with python ref executor tests, none of the + # data dependent pytorch ops actually decompose to prims + tags_intersection.discard(torch.Tag.data_dependent_output) + + # iter over first tags for determinism + _prim._tags = tuple(t for t in overload_tags[0] if t in tags_intersection) + + from torch._subclasses.fake_tensor import contains_tensor_types + + if not any(contains_tensor_types(a.type) for a in _prim._schema.arguments) or str( + _prim + ) in [ + # See https://github.com/pytorch/pytorch/issues/103532 + "prims.device_put.default" + ]: + prim_backend_select_impl.impl(name, _backend_select_impl) + + for p in (_prim_packet, _prim): + p.__doc__ = doc + p.return_type = return_type # type: ignore[attr-defined] + + p.schema = schema + p.prim_impl = _prim_impl + p.prim_meta_impl = meta + p.impl_aten = impl_aten + + return _prim + + +class ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND(Enum): + DEFAULT = (0,) + INT_TO_FLOAT = (2,) + ALWAYS_BOOL = (3,) + COMPLEX_TO_FLOAT = (4,) + + +# TODO: implement dtype validation here, too, or on the corresponding refs +def _prim_elementwise_meta( + *args, + type_promotion: ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND, + args_with_fixed_dtypes: Optional[Tuple[TensorLikeType, ...]] = None, +) -> FakeTensor: + """ + Meta function for elementwise operations that produce outputs in the same dtype + as their inputs. + + Stride logic is currently incorrect. + """ + + assert len(args) > 0 + + utils.check_same_dtype(*args) + + args_ = list(args) + if args_with_fixed_dtypes is not None: + args_ = list(args_with_fixed_dtypes) + args_ + + utils.check_same_device(*args_, allow_cpu_scalar_tensors=True) + utils.check_same_shape(*args_, allow_cpu_scalar_tensors=True) + + l2p_perm = utils.compute_elementwise_output_logical_to_physical_perm(*args_) + shape = utils.extract_shape(*args_, allow_cpu_scalar_tensors=True) + + # Acquires the dtype + dtype = None + scalar_type = None + for arg in args: + if isinstance(arg, TensorLike): + if not utils.is_cpu_scalar_tensor(arg): + dtype = arg.dtype + break + else: + dtype = arg.dtype + elif isinstance(arg, Number): + scalar_type = type(arg) + + if dtype is None and scalar_type is not None: + dtype = utils.type_to_dtype(scalar_type) + + # Acquires the device (if it exists) or number + device = None + number = None + for arg in args_: + if isinstance(arg, TensorLike): + if utils.is_cpu_scalar_tensor(arg): + if device is None: + device = arg.device + # keep going, in case there is a cuda tensor later + else: + device = arg.device + break + + elif isinstance(arg, Number): + if number is None: + number = arg + + # NOTE: type promotion behavior here is mostly hidden from tests because + # references will typically handle the type promotion properly even if this doesn't + # (but getting it wrong will cause too many casts to be inserted in traces!) + if device is not None: + assert dtype is not None + if type_promotion == ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT: + dtype = dtype + elif type_promotion == ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.ALWAYS_BOOL: + dtype = torch.bool + elif type_promotion == ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.INT_TO_FLOAT: + if utils.is_integer_dtype(dtype) or utils.is_boolean_dtype(dtype): + dtype = torch.get_default_dtype() + elif type_promotion == ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT: + if utils.is_complex_dtype(dtype): + dtype = utils.corresponding_real_dtype(dtype) + else: + dtype = dtype + + assert shape is not None + return torch.empty_permuted(shape, l2p_perm, device=device, dtype=dtype) # type: ignore[return-value] + + # Number case + # TODO: fix number type promotion (bool, complex->float) + + # For now for symint/float, just implementing the common / simple cases of (int,float,symint,symfloat) + seen_float = False + if isinstance(number, (torch.SymInt, torch.SymFloat)): + for a in args: + assert isinstance(a, (int, float, torch.SymInt, torch.SymFloat)), "NYI" + seen_float = seen_float or isinstance(a, (float, torch.SymFloat)) + if seen_float: + number = sym_float(number) + + return TensorMeta(number) # type: ignore[arg-type] + + +def _complex_only_elementwise_meta(*args, **kwargs): + torch._check( + utils.is_complex_dtype(args[0].dtype), lambda: "Only complex dtype is supported" + ) + return _prim_elementwise_meta(*args, **kwargs) + + +def _make_elementwise_unary_prim( + name: str, *, type_promotion: ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND, **kwargs +): + """ + Creates an elementwise unary prim. + """ + + return _make_prim( + schema=f"{name}(Tensor self) -> Tensor", + meta=partial(_prim_elementwise_meta, type_promotion=type_promotion), + return_type=RETURN_TYPE.NEW, + **kwargs, + ) + + +def _make_elementwise_binary_prim( + name: str, *, type_promotion: ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND, **kwargs +): + """ + Creates an elementwise binary prim. + """ + + return _make_prim( + schema=f"{name}(Tensor self, Tensor other) -> Tensor", + meta=partial(_prim_elementwise_meta, type_promotion=type_promotion), + return_type=RETURN_TYPE.NEW, + **kwargs, + ) + + +def _not_impl(*args, **kwargs): + raise NotImplementedError + + +# +# Elementwise unary operations +# + + +abs = _make_elementwise_unary_prim( + "abs", + impl_aten=torch.abs, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT, +) + +acos = _make_elementwise_unary_prim( + "acos", + impl_aten=torch.acos, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +acosh = _make_elementwise_unary_prim( + "acosh", + impl_aten=torch.acosh, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +asin = _make_elementwise_unary_prim( + "asin", + impl_aten=torch.asin, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +asinh = _make_elementwise_unary_prim( + "asinh", + impl_aten=torch.asinh, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +atan = _make_elementwise_unary_prim( + "atan", + impl_aten=torch.atan, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +atanh = _make_elementwise_unary_prim( + "atanh", + impl_aten=torch.atanh, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +cos = _make_elementwise_unary_prim( + "cos", + impl_aten=torch.cos, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +cosh = _make_elementwise_unary_prim( + "cosh", + impl_aten=torch.cosh, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bessel_j0 = _make_elementwise_unary_prim( + "bessel_j0", + impl_aten=torch.special.bessel_j0, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bessel_j1 = _make_elementwise_unary_prim( + "bessel_j1", + impl_aten=torch.special.bessel_j1, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bessel_i0 = _make_elementwise_unary_prim( + "bessel_i0", + impl_aten=torch.i0, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bessel_i0e = _make_elementwise_unary_prim( + "bessel_i0e", + impl_aten=torch.special.i0e, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bessel_i1 = _make_elementwise_unary_prim( + "bessel_i1", + impl_aten=torch.special.i1, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bessel_i1e = _make_elementwise_unary_prim( + "bessel_i1e", + impl_aten=torch.special.i1e, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bitwise_not = _make_elementwise_unary_prim( + "bitwise_not", + impl_aten=torch.bitwise_not, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + + +def _cbrt_aten(a: torch.Tensor) -> Tensor: + torch._check( + not a.is_complex(), + lambda: "cbrt: Complex inputs not supported. Consider calling torch.pow(a, 1.0/3.0)", + ) + # Returns the real cubic root of the number. + # Note that if a < 0, pow(a, (1. / 3.)) returns th complex number + # exp(1/3 * log(a)) = exp(1/3 * (log(abs(a)) + pi*i)) = cbrt(abs(a)) * e^{pi/3*i} + # which is a complex number. + # For more info see the section Note in + # https://en.cppreference.com/w/cpp/numeric/math/cbrt + return torch.copysign(torch.pow(a.abs(), 1 / 3), a) + + +cbrt = _make_elementwise_unary_prim( + "cbrt", + impl_aten=_cbrt_aten, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +ceil = _make_elementwise_unary_prim( + "ceil", + impl_aten=torch.ceil, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + + +def _conj_physical_meta(input: TensorLikeType) -> TensorLikeType: + if not input.dtype.is_complex: + raise RuntimeError("prims.conj_physical is only defined for complex dtypes") + + strides = utils.compute_elementwise_output_strides(input) + return TensorMeta(input, strides=strides) + + +conj_physical = _make_prim( + schema="conj_physical(Tensor self) -> Tensor", + meta=_conj_physical_meta, + impl_aten=torch._conj_physical, + doc="Returns the physical conjugation of a complex tensor", + return_type=RETURN_TYPE.NEW, +) + + +def _clone_meta( + input: TensorLikeType, *, memory_format: torch.memory_format = torch.preserve_format +) -> TensorLikeType: + if memory_format != torch.preserve_format: + return torch.empty( + input.shape, + dtype=input.dtype, + layout=input.layout, + device=input.device, + memory_format=memory_format, + ) + + # memory_format == torch.preserve_format + strides = utils.compute_elementwise_output_strides(input) + return torch.empty_strided( + input.shape, + strides, + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + +clone = _make_prim( + schema="clone(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor", + meta=_clone_meta, + impl_aten=torch.clone, + doc="Returns the copy of a tensor", + return_type=RETURN_TYPE.NEW, + register_conj_neg_fallthrough=True, +) + +digamma = _make_elementwise_unary_prim( + "digamma", + impl_aten=torch.digamma, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +erf = _make_elementwise_unary_prim( + "erf", + impl_aten=torch.erf, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +erf_inv = _make_elementwise_unary_prim( + "erf_inv", + impl_aten=torch.special.erfinv, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +erfc = _make_elementwise_unary_prim( + "erfc", + impl_aten=torch.special.erfc, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +erfcx = _make_elementwise_unary_prim( + "erfcx", + impl_aten=torch.special.erfcx, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +exp = _make_elementwise_unary_prim( + "exp", + impl_aten=torch.exp, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +expm1 = _make_elementwise_unary_prim( + "expm1", + impl_aten=torch.special.expm1, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +exp2 = _make_elementwise_unary_prim( + "exp2", + impl_aten=torch.special.exp2, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + + +def _fill_meta(a: TensorLikeType, value: NumberType) -> TensorLikeType: + return _prim_elementwise_meta( + a, type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT + ) + + +# NOTE: fill uses _make_prim directly because it has a value parameter +fill = _make_prim( + schema="fill(Tensor self, Scalar value) -> Tensor", + return_type=RETURN_TYPE.NEW, + meta=_fill_meta, + impl_aten=torch.fill, + doc="", +) + +floor = _make_elementwise_unary_prim( + "floor", + impl_aten=torch.floor, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +imag = _make_prim( + schema="imag(Tensor(a) self) -> Tensor(a)", + meta=partial( + _complex_only_elementwise_meta, + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT, + ), + return_type=RETURN_TYPE.VIEW, + impl_aten=torch.imag, + doc="", +) + +isfinite = _make_elementwise_unary_prim( + "isfinite", + impl_aten=torch.isfinite, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.ALWAYS_BOOL, +) + +lgamma = _make_elementwise_unary_prim( + "lgamma", + impl_aten=torch.lgamma, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +log = _make_elementwise_unary_prim( + "log", + impl_aten=torch.log, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +log1p = _make_elementwise_unary_prim( + "log1p", + impl_aten=torch.log1p, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +log2 = _make_elementwise_unary_prim( + "log2", + impl_aten=torch.log2, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +log10 = _make_elementwise_unary_prim( + "log10", + impl_aten=torch.log10, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +real = _make_prim( + schema="real(Tensor(a) self) -> Tensor(a)", + meta=partial( + _complex_only_elementwise_meta, + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT, + ), + return_type=RETURN_TYPE.VIEW, + impl_aten=torch.real, + doc="", +) + +reciprocal = _make_elementwise_unary_prim( + "reciprocal", + impl_aten=torch.reciprocal, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +ndtri = _make_elementwise_unary_prim( + "ndtri", + impl_aten=torch.special.ndtri, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +neg = _make_elementwise_unary_prim( + "neg", + impl_aten=torch.neg, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +round = _make_elementwise_unary_prim( + "round", + impl_aten=torch.round, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +rsqrt = _make_elementwise_unary_prim( + "rsqrt", + impl_aten=torch.rsqrt, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +sign = _make_elementwise_unary_prim( + "sign", + impl_aten=torch.sign, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +signbit = _make_elementwise_unary_prim( + "signbit", + impl_aten=torch.signbit, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +sin = _make_elementwise_unary_prim( + "sin", + impl_aten=torch.sin, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +sinh = _make_elementwise_unary_prim( + "sinh", + impl_aten=torch.sinh, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +spherical_bessel_j0 = _make_elementwise_unary_prim( + "spherical_bessel_j0", + impl_aten=torch.special.spherical_bessel_j0, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +sqrt = _make_elementwise_unary_prim( + "sqrt", + impl_aten=torch.sqrt, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +tan = _make_elementwise_unary_prim( + "tan", + impl_aten=torch.tan, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +tanh = _make_elementwise_unary_prim( + "tanh", + impl_aten=torch.tanh, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +trunc = _make_elementwise_unary_prim( + "trunc", + impl_aten=torch.trunc, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +# +# Elementwise binary operations +# + +add = _make_elementwise_binary_prim( + name="add", + impl_aten=torch.add, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +atan2 = _make_elementwise_binary_prim( + name="atan2", + impl_aten=torch.atan2, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bitwise_and = _make_elementwise_binary_prim( + "bitwise_and", + impl_aten=torch.bitwise_and, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bitwise_or = _make_elementwise_binary_prim( + "bitwise_or", + impl_aten=torch.bitwise_or, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +bitwise_xor = _make_elementwise_binary_prim( + "bitwise_xor", + impl_aten=torch.bitwise_xor, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +# TODO: complex needs a special meta to account for its float -> complex behavior +# complex = _make_elementwise_binary_prim( +# impl_aten=torch.complex, +# doc="", +# ) + + +# div prim performs truncation division on integer inputs +# and true division for floating and complex inputs +def _div_aten(a, b): + is_integral = isinstance(a, (bool, int, torch.SymInt)) or ( + isinstance(a, torch.Tensor) and utils.is_integer_dtype(a.dtype) + ) + + if is_integral: + return torch.div(a, b, rounding_mode="trunc") + else: + return torch.true_divide(a, b) + + +div = _make_elementwise_binary_prim( + "div", + impl_aten=_div_aten, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +eq = _make_elementwise_binary_prim( + "eq", + impl_aten=torch.eq, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.ALWAYS_BOOL, +) + +fmax = _make_elementwise_binary_prim( + "fmax", + impl_aten=torch.fmax, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +fmin = _make_elementwise_binary_prim( + "fmin", + impl_aten=torch.fmin, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +fmod = _make_elementwise_binary_prim( + "fmod", + impl_aten=torch.fmod, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + + +gcd = _make_elementwise_binary_prim( + "gcd", + impl_aten=torch.gcd, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + + +ge = _make_elementwise_binary_prim( + "ge", + impl_aten=torch.ge, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.ALWAYS_BOOL, +) + +gt = _make_elementwise_binary_prim( + "gt", + impl_aten=torch.gt, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.ALWAYS_BOOL, +) + +hypot = _make_elementwise_binary_prim( + "hypot", + impl_aten=torch.hypot, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +igamma = _make_elementwise_binary_prim( + "igamma", + impl_aten=torch.special.gammainc, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +igammac = _make_elementwise_binary_prim( + "igammac", + impl_aten=torch.special.gammaincc, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +le = _make_elementwise_binary_prim( + "le", + impl_aten=torch.le, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.ALWAYS_BOOL, +) + +lt = _make_elementwise_binary_prim( + "lt", + impl_aten=torch.lt, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.ALWAYS_BOOL, +) + + +# Note: the following impls are because torch.maximum and torch.minimum do not support scalar inputs +def _maximum_aten( + a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberType] +) -> TensorLikeType: + if isinstance(a, TensorLike) and isinstance(b, Number): + b = scalar_tensor(b, dtype=a.dtype, device=a.device) + elif isinstance(b, TensorLike) and isinstance(a, Number): + a = scalar_tensor(a, dtype=b.dtype, device=b.device) + + return torch.maximum(a, b) # type: ignore[arg-type] + + +maximum = _make_elementwise_binary_prim( + "maximum", + impl_aten=_maximum_aten, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + + +def _minimum_aten( + a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberType] +) -> TensorLikeType: + if isinstance(a, TensorLike) and isinstance(b, Number): + b = scalar_tensor(b, dtype=a.dtype, device=a.device) + elif isinstance(b, TensorLike) and isinstance(a, Number): + a = scalar_tensor(a, dtype=b.dtype, device=b.device) + + return torch.minimum(a, b) # type: ignore[arg-type] + + +minimum = _make_elementwise_binary_prim( + "minimum", + impl_aten=_minimum_aten, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +mul = _make_elementwise_binary_prim( + "mul", + impl_aten=torch.mul, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +ne = _make_elementwise_binary_prim( + "ne", + impl_aten=torch.ne, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.ALWAYS_BOOL, +) + +nextafter = _make_elementwise_binary_prim( + "nextafter", + impl_aten=torch.nextafter, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +pow = _make_elementwise_binary_prim( + "pow", + impl_aten=torch.pow, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +remainder = _make_elementwise_binary_prim( + "remainder", + impl_aten=torch.remainder, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + + +shift_left = _make_elementwise_binary_prim( + "shift_left", + impl_aten=torch.bitwise_left_shift, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +shift_right_arithmetic = _make_elementwise_binary_prim( + "shift_right_arithmetic", + impl_aten=torch.bitwise_right_shift, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +shift_right_logical = _not_impl + +sub = _make_elementwise_binary_prim( + "sub", + impl_aten=torch.sub, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + +zeta = _make_elementwise_binary_prim( + "zeta", + impl_aten=torch.special.zeta, + doc="", + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, +) + + +# +# View operations +def _as_strided_meta( + a: TensorLikeType, size: ShapeType, stride: StrideType, storage_offset: int +) -> TensorLikeType: + assert len(size) == len(stride) + assert storage_offset >= 0 + utils.validate_strides(stride) + utils.validate_shape(size) + + if reduce(operator.mul, size) == 0: + # NOTE: This special case is to avoid having to acquire the storage below + # as_strided to shapes with no elements are trivially valid, so it's OK + pass + elif isinstance(a, torch.Tensor): + utils.check_in_bounds_for_storage( + a._typed_storage(), size, stride, storage_offset + ) + + return torch.as_strided(a, size, stride, storage_offset) + + +def _as_strided_aten( + a: Tensor, size: ShapeType, stride: StrideType, storage_offset: int +) -> Tensor: + return torch.as_strided(a, size, stride, storage_offset) + + +_as_strided_doc = """ + Creates a view of the tensor with the given shape (size), strides (stride) and + storage offset (storage_offset). +""" + +as_strided = _make_prim( + schema="as_strided(Tensor(a!) a, SymInt[] size, SymInt[] stride, SymInt storage_offset) -> Tensor(a!)", + meta=_as_strided_meta, + impl_aten=_as_strided_aten, + return_type=RETURN_TYPE.VIEW, + doc=_as_strided_doc, +) + + +def _broadcast_in_dim_meta( + a: TensorLikeType, shape: ShapeType, broadcast_dimensions: Sequence[int] +): + from torch.fx.experimental.symbolic_shapes import guard_size_oblivious + + # Type checks + assert isinstance(a, TensorLike) + assert isinstance(shape, Sequence) + assert isinstance(broadcast_dimensions, Sequence) + + # every dimension must be accounted for + assert a.ndim == len(broadcast_dimensions) + + # broadcast shape must have weakly more dimensions + assert len(shape) >= a.ndim + + # broadcast_dimensions must be an ascending sequence + # (no relative reordering of dims) of integers and + # each dimension must be within the new shape + def _greater_than_reduce(acc, x): + assert isinstance(x, Dim) + assert x > acc + assert x < len(shape) + + return x + + reduce(_greater_than_reduce, broadcast_dimensions, -1) + + # shape must be broadcastable to + for idx, new_idx in enumerate(broadcast_dimensions): + if not guard_size_oblivious(a.shape[idx] == 1): + torch._check( + a.shape[idx] == shape[new_idx], + lambda: f"{a.shape[idx]} must be broadcastable to {shape[new_idx]}", + ) + + new_strides = [] + original_idx = 0 + for idx in range(len(shape)): + if idx in broadcast_dimensions: + # Assigns a stride of zero to dimensions + # which were actually broadcast + if guard_size_oblivious(a.shape[original_idx] != shape[idx]): + new_strides.append(0) + else: + new_strides.append(a.stride()[original_idx]) + original_idx = original_idx + 1 + else: + if guard_size_oblivious(shape[idx] != 1): + new_strides.append(0) + elif original_idx == a.ndim: + new_strides.append(1) + else: + new_strides.append(a.stride()[original_idx] * a.size()[original_idx]) + + return a.as_strided(shape, new_strides, a.storage_offset()) + + +def _broadcast_in_dim_aten(a, shape, broadcast_dimensions): + s = list(shape) + for broadcast_dimension in broadcast_dimensions: + s[broadcast_dimension] = -1 + + v = a + for idx, x in enumerate(s): + if x != -1: + v = v.unsqueeze(idx) + + return v.expand(shape) + + +_broadcast_in_dim_doc = """ + Creates a view of a with the specified shape. + + Allows adding dimensions of any length and broadcasting + dimensions of length one in a to any length. + + The location of the broadcast dimensions must be specified + using the broadcast_dimensions argument. Changing the + relative order of dimensions is not supported. + """ + +broadcast_in_dim = _make_prim( + schema="broadcast_in_dim(Tensor(a) a, SymInt[] shape, int[] broadcast_dimensions) -> Tensor(a)", + meta=_broadcast_in_dim_meta, + impl_aten=_broadcast_in_dim_aten, + return_type=RETURN_TYPE.VIEW, + doc=_broadcast_in_dim_doc, +) + + +def _validate_collapse_args(a: Tensor, start: int, end: int) -> None: + # Special-case for zero dimensional tensors + ndim = max(1, a.dim()) + utils.validate_idx(ndim, start) + utils.validate_idx(ndim, end) + + # Verifies end is strictly greater than start + # (Collapse requires a non-empty interval) + torch._check_value( + end >= start, + lambda: f"Attempting to collapse but end, {end}, is less than start, {start}!", + ) + + +def _collapsed_shape(shape: ShapeType, start: int, end: int) -> Tuple[int, ...]: + """ + Returns the shape of a with dims in [start, end) merged into a single dimension. + """ + # Special-case for zero dimensional tensors + shape = (1,) if len(shape) == 0 else tuple(shape) + + dim_length = 1 + for s in shape[start : end + 1]: + dim_length = dim_length * s + + return shape[0:start] + (dim_length,) + shape[end + 1 :] + + +def _collapse_view_helper( + a: TensorLikeType, start: int, end: int +) -> Tuple[Optional[ShapeType], Optional[StrideType]]: + assert isinstance(a, TensorLike) + + from torch.fx.experimental.symbolic_shapes import guard_size_oblivious + + _validate_collapse_args(a, start, end) + + # Special-case for zero dimensional tensors + if a.ndim == 0: + shape = (1,) + strides = (1,) + else: + shape = a.shape # type: ignore[assignment] + strides = a.stride() # type: ignore[assignment] + + if a.ndim == 0 or (end == start): + return shape, strides + + length = shape[end] + stride = strides[end] + for idx in range(end - 1, start - 1, -1): + if guard_size_oblivious(shape[idx] == 0) or guard_size_oblivious( + shape[idx + 1] == 0 + ): + length = 0 + stride = 0 + break + + if guard_size_oblivious(shape[idx] == 1): + continue + + length = length * shape[idx] + if guard_size_oblivious(stride < strides[idx]): + stride = stride + else: + stride = strides[idx] + + if ( + guard_size_oblivious(a.numel() > 0) + and guard_size_oblivious(shape[idx + 1] != 1) + and not guard_size_oblivious( + strides[idx] == strides[idx + 1] * shape[idx + 1] + ) + ): + return None, None + + new_shape = shape[:start] + (length,) + shape[end + 1 :] + new_strides = strides[:start] + (stride,) + strides[end + 1 :] + + # NOTE: when the input has no elements it's restrided as if it were contiguous + if guard_size_oblivious(a.numel() == 0): + new_strides = utils.make_contiguous_strides_for(new_shape) + + return new_shape, new_strides + + +def _collapse_view_meta(a: TensorLikeType, start: int, end: int) -> TensorLikeType: + new_shape, new_strides = _collapse_view_helper(a, start, end) + + if new_shape is None: + msg = "Attempting to view a collapsed tensor, but no such view exists!" + raise ValueError(msg) + + assert new_strides is not None + return a.as_strided(new_shape, new_strides, a.storage_offset()) + + +def _collapse_view_aten(a: Tensor, start: int, end: int) -> Tensor: + new_shape = _collapsed_shape(a.shape, start, end) + return a.view(new_shape) + + +_collapse_view_doc = """ + Creates a view of a with the dimensions between + start (inclusive) and end (exclusive) merged into a + single dimension. + + If it's not possible to take such a view then an error + is thrown. See collapse instead. + + The dimensions can be merged if and only if + they are all "nested" with each other. That is, they all + have the property that + + stride[i] = stride[i+1] * shape[i+1] + + for all i in [start, end - 1). + """ + +collapse_view = _make_prim( + schema="collapse_view(Tensor(a) a, int start, int end) -> Tensor(a)", + meta=_collapse_view_meta, + impl_aten=_collapse_view_aten, + return_type=RETURN_TYPE.VIEW, + doc=_collapse_view_doc, +) + + +def _conj_meta(a: TensorLikeType) -> TensorLikeType: + if not a.dtype.is_complex: + raise RuntimeError("Expected complex dtype in prims.conj") + out = a.as_strided(a.shape, a.stride(), a.storage_offset()) + torch._C._set_conj(out, not a.is_conj()) + return out + + +_conj_doc = """ +Returns a conjugated view of the original tensor +""" + +conj = _make_prim( + schema="conj(Tensor(a) a) -> Tensor(a)", + meta=_conj_meta, + impl_aten=torch.conj, + return_type=RETURN_TYPE.VIEW, + doc=_conj_doc, +) + + +def expand_dims( + a: TensorLikeType, dimensions: DimsSequenceType, ndim=None +) -> TensorLikeType: + """ + Creates a view of a with a.ndim + len(dimensions) dimensions, with new + dimensions of length one at the dimensions specified by dimensions. + """ + if ndim is not None: + # TODO: this is only here to support the unsqueeze ref + dims = sorted(utils.canonicalize_dims(ndim, dimensions)) # type: ignore[arg-type] + else: + dims = sorted(utils.canonicalize_dims(a.ndim, dimensions)) # type: ignore[arg-type] + if len(set(dims)) != len(dims): + msg = f"Received duplicate dimensions to expand in {str(dimensions)}" + raise ValueError(msg) + + new_shape = list(a.shape) + for idx in dims: + new_shape.insert(idx, 1) + + broadcast_dimensions = [ + idx for idx in range(len(new_shape)) if idx not in dimensions + ] + return broadcast_in_dim(a, new_shape, broadcast_dimensions) + + +def _split_dim_meta(a: TensorLikeType, dim: int, outer_length: int) -> TensorLikeType: + assert isinstance(a, TensorLike) + utils.validate_idx(a.ndim, dim) + utils.validate_dim_length(outer_length) + + # Verifies the dim can be split with the specified lhs_length + inner_length = a.shape[dim] // outer_length + + if (a.shape[dim] % outer_length) != 0: + msg = ( + f"Attempting to split dimension of length {a.shape[dim]}, " + f"but outer length of {outer_length} divides it with a remainder!" + ) + raise ValueError(msg) + + new_shape: List[int] = [] + new_strides: List[int] = [] + for idx in range(a.ndim): + if idx == dim: + new_shape.extend((outer_length, inner_length)) + new_strides.extend((a.stride()[idx] * inner_length, a.stride()[idx])) + else: + new_shape.append(a.shape[idx]) + new_strides.append(a.stride()[idx]) + + return a.as_strided(new_shape, new_strides, a.storage_offset()) + + +def _split_dim_aten(a: Tensor, dim: int, outer_length: int) -> Tensor: + inner_length = a.shape[dim] // outer_length + new_shape = a.shape[0:dim] + (outer_length, inner_length) + a.shape[dim + 1 :] + + return a.view(new_shape) + + +_split_dim_doc = """ + Creates a view of a with the given dimension (of length l) split + into two dimensions, with the outer of the two having + length outer_length and the inner of the two having computed + length inner_length such outer_length * inner_length = l. + """ + +# TODO: consider renaming split_dim_view +split_dim = _make_prim( + schema="split_dim(Tensor(a) a, int dim, SymInt outer_length) -> Tensor(a)", + meta=_split_dim_meta, + impl_aten=_split_dim_aten, + return_type=RETURN_TYPE.VIEW, + doc=_split_dim_doc, +) + + +# Note: allows dimensions to be specified redundantly +def _squeeze_meta(a: TensorLikeType, dimensions: Sequence) -> TensorLikeType: + assert isinstance(a, TensorLike) + + for idx in dimensions: + utils.validate_idx(a.ndim, idx) + assert a.shape[idx] == 1 + + new_shape = [] + new_strides = [] + for idx in range(len(a.shape)): + if idx in dimensions: + continue + + new_shape.append(a.shape[idx]) + new_strides.append(a.stride()[idx]) + + return a.as_strided(new_shape, new_strides, a.storage_offset()) + + +_squeeze_doc = """ + Creates a view of the tensor with the specified dimensions removed. + + The removed dimensions must each have length one. + """ + +squeeze = _make_prim( + schema="squeeze(Tensor(a) a, int[] dimensions) -> Tensor(a)", + meta=_squeeze_meta, + impl_aten=torch.squeeze, + return_type=RETURN_TYPE.VIEW, + doc=_squeeze_doc, +) + + +def _transpose_meta(a: TensorLikeType, permutation: DimsSequenceType) -> TensorLikeType: + if a.ndim != len(permutation): + msg = f"Attempting to permute a tensor of rank {a.ndim}, but received a permutation of length {len(permutation)}!" + raise ValueError(msg) + + if not utils.is_valid_permutation(a.ndim, permutation): + msg = f"Received an invalid permutation, {permutation}!" + raise ValueError(msg) + + new_shape = [0] * a.ndim + new_strides = [0] * a.ndim + for idx, dim in enumerate(permutation): + new_shape[idx] = a.shape[dim] + new_strides[idx] = a.stride()[dim] + + return a.as_strided(tuple(new_shape), tuple(new_strides), a.storage_offset()) + + +def _transpose_aten(a: Tensor, permutation: DimsSequenceType) -> Tensor: + return torch.permute(a, permutation) + + +_transpose_doc = """ + Creates a view of the tensor with its dimensions permuted. + + The length of the permutation must be the rank of the tensor, + and each element of the permutation specifies the new order + for the corresponding dimension. + """ + +transpose = _make_prim( + schema="transpose(Tensor(a) a, int[] permutation) -> Tensor(a)", + meta=_transpose_meta, + impl_aten=_transpose_aten, + return_type=RETURN_TYPE.VIEW, + doc=_transpose_doc, +) + + +def _view_of_meta(a: TensorLikeType) -> TensorLikeType: + return a.as_strided(a.shape, a.stride(), a.storage_offset()) + + +def _view_of_aten(a: Tensor) -> Tensor: + return a.view(a.shape) + + +_view_of_doc = """ + Creates a view of the tensor. + """ + +view_of = _make_prim( + schema="view_of(Tensor(a) a) -> Tensor(a)", + meta=_view_of_meta, + impl_aten=_view_of_aten, + return_type=RETURN_TYPE.VIEW, + doc=_view_of_doc, +) + + +def _view_element_type_meta(a: TensorLikeType, dtype: torch.dtype) -> TensorLikeType: + return a.view(dtype) + + +def _view_element_type_aten(a: Tensor, dtype: torch.dtype) -> Tensor: + return a.view(dtype) + + +_view_element_type_doc = """ + Creates a view of the tensor with a different dtype. + """ + +view_element_type = _make_prim( + schema="view_of_dtype(Tensor(a) a, ScalarType dtype) -> Tensor(a)", + meta=_view_element_type_meta, + impl_aten=_view_element_type_aten, + return_type=RETURN_TYPE.VIEW, + doc=_view_element_type_doc, +) + +# +# Functionalized view mutations +# + + +def _as_strided_scatter_meta( + input: TensorLikeType, + src: TensorLikeType, + size: ShapeType, + stride: StrideType, + storage_offset: int, +) -> TensorLikeType: + utils.validate_shape(size) + utils.validate_strides(stride) + + required_size = utils.compute_required_storage_length(size, stride, storage_offset) + torch._check( + input.numel() >= required_size, + lambda: ( + f"as_strided_scatter: sizes {size}, strides {stride}, storage offset {storage_offset} " + f" and itemsize {input.element_size()} requiring a storage size of " + f"{required_size * input.element_size()} are out of bounds " + f"for storage of size {input.numel() * input.element_size()}" + ), + ) + torch._check( + utils.is_same_shape(src.shape, size), + lambda: f"expected src to have a size equal to the slice of self. src size = {src.shape}, slice size = {size}", + ) + + return utils.clone_preserve_strides(input) + + +_as_strided_scatter_doc = """ + Creates a new tensor equivalent to ``out = input.clone()`` after mutation by + ``out.as_strided(size, stride, storage_offset).copy_(src)``. +""" + +as_strided_scatter = _make_prim( + schema="as_strided_scatter(Tensor self, Tensor src, SymInt[] size, SymInt[] stride, SymInt storage_offset) -> Tensor", + meta=_as_strided_scatter_meta, + impl_aten=torch.as_strided_scatter, + return_type=RETURN_TYPE.NEW, + doc=_as_strided_scatter_doc, +) + + +# +# Shape operations +# + + +def _collapse_meta(a: Tensor, start: int, end: int) -> Tensor: + # Special-case for zero dimensional tensors + _validate_collapse_args(a, start, end) + new_shape = _collapsed_shape(a.shape, start, end) + return a.new_empty(new_shape) + + +def _collapse_aten(a: Tensor, start: int, end: int) -> Tensor: + new_shape = _collapsed_shape(a.shape, start, end) + out = a.new_empty(new_shape) + with torch.no_grad(): + out.view_as(a).copy_(a) + return out + + +_collapse_doc = """ +Collapse a span of neighboring dimensions into one. + +See collapse_view for the corresponding view operation. +""" +collapse = _make_prim( + schema="collapse(Tensor a, int start, int end) -> Tensor", + meta=_collapse_meta, + impl_aten=_collapse_aten, + return_type=RETURN_TYPE.NEW, + doc=_collapse_doc, +) + + +# TODO: review stride logic +# NB: unlike torch.cat, this is more strict about empty tensors and dim is +# never negative +def _cat_meta(tensors: Sequence[TensorLikeType], dim: int) -> TensorLikeType: + # Verifies same shape (except in the concat dimension) + assert dim >= 0 + shape = tensors[0].shape + sym_sum_args = [] + for tensor_idx, tensor in enumerate(tensors): + assert len(shape) == len(tensor.shape) + for idx, (common_length, length) in enumerate(zip(shape, tensor.shape)): + if idx == dim: + sym_sum_args.append(length) + else: + torch._check( + length == common_length, + lambda: f"Sizes of tensors must match except in dimension {dim}. " + f"Expected {common_length} in dimension {idx} but got {length} for tensor number " + f"{tensor_idx} in the list", + ) + + new_shape = list(tensors[0].shape).copy() + new_shape[dim] = torch.sym_sum(sym_sum_args) + return TensorMeta( + tensors[0], + shape=new_shape, + strides=utils.make_contiguous_strides_for(new_shape), + ) + + +def _cat_aten(tensors: Union[Tuple[Tensor, ...], List[Tensor]], dim: int) -> Tensor: + return torch.cat(tensors, dim) + + +_cat_doc = """ + Concatenates tensors along the specified dimension. + + The tensors' shapes must have the same rank and same length for other dimensions. + """ + +cat = _make_prim( + schema="cat(Tensor[] tensors, int dim) -> Tensor", + meta=_cat_meta, + impl_aten=_cat_aten, + return_type=RETURN_TYPE.NEW, + doc=_cat_doc, +) + + +def _reshape_meta(a: TensorLikeType, shape: ShapeType): + assert isinstance(a, TensorLike) + utils.validate_shape(shape) + + # Validates the tensor and the requested shape have the + # same number of elements + numel = reduce(operator.mul, shape) + if numel != a.numel(): + msg = f"Attempting to reshape a tensor with {a.numel()} elements to a shape with {numel} elements!" + raise ValueError(msg) + + return TensorMeta(a, shape=shape, strides=utils.make_contiguous_strides_for(shape)) + + +def _reshape_aten(a: Tensor, shape: ShapeType) -> Tensor: + return a.reshape(shape).contiguous().clone() + + +_reshape_doc = """ + Creates a contiguous tensor with the specified shape + containing a copy of the data in a. + """ +reshape = _make_prim( + schema="reshape(Tensor a, SymInt[] shape) -> Tensor", + meta=_reshape_meta, + impl_aten=_reshape_aten, + return_type=RETURN_TYPE.NEW, + doc=_reshape_doc, +) + + +def _rev_meta(a: TensorLikeType, dims: DimsSequenceType) -> TensorLikeType: + utils.validate_dimension_indices(a.ndim, dims) + return torch.empty_like(a, memory_format=torch.preserve_format) + + +_rev_doc = """ + Reverses the order of elements along the given dimensions. + """ + +rev = _make_prim( + schema="rev(Tensor a, int[] dims) -> Tensor", + meta=_rev_meta, + impl_aten=torch.flip, + return_type=RETURN_TYPE.NEW, + doc=_rev_doc, +) + +# +# Conditional prims +# + + +def _where_meta( + pred: TensorLikeType, a: TensorLikeType, b: TensorLikeType +) -> TensorLikeType: + return _prim_elementwise_meta( + a, + b, + type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT, + args_with_fixed_dtypes=(pred,), + ) + + +_where_doc = """ + Selects elements from a and b according to pred. + + Where pred is true the result contains the element from a, and + where pred is false the result contains the element from b. + """ + +where = _make_prim( + schema="where(Tensor pred, Tensor a, Tensor b) -> Tensor", + meta=_where_meta, + impl_aten=torch.where, + return_type=RETURN_TYPE.NEW, + doc=_where_doc, +) + + +# +# Type conversions +# +def _convert_element_type_meta(a: TensorLikeType, dtype: torch.dtype) -> TensorLikeType: + # Type checks + assert isinstance(a, TensorLike) + assert isinstance(dtype, torch.dtype) + + # dtype conversion preserves dense strides + if torch._prims_common.is_non_overlapping_and_dense(a): + strides = a.stride() + else: + strides = utils.compute_elementwise_output_strides(a) + + return TensorMeta(a, strides=strides, dtype=dtype) + + +def _convert_element_type_aten(a: Tensor, dtype: torch.dtype) -> Tensor: + # Propagates requires grad when possible + if not utils.is_grad_dtype(dtype): + requires_grad = False + else: + # TODO: update meta objects so this can be acquired directly + try: + requires_grad = a.requires_grad + except Exception as e: + requires_grad = False + + result = torch.empty_like( + a, device=a.device, dtype=dtype, requires_grad=requires_grad + ) + with torch.no_grad(): + return copy_to(result, a) + + +_convert_element_type_doc = """ + Creates a copy of a tensor with the given dtype. + """ + +convert_element_type = _make_prim( + schema="convert_element_type(Tensor a, ScalarType dtype) -> Tensor", + meta=_convert_element_type_meta, + impl_aten=_convert_element_type_aten, + return_type=RETURN_TYPE.NEW, + doc=_convert_element_type_doc, + tags=(torch.Tag.pointwise,), +) + + +def _device_put_meta( + a: TensorLikeType, device: Union[str, torch.device], non_blocking=False +) -> TensorLikeType: + assert isinstance(a, TensorLike) + assert isinstance(device, (str, torch.device)) + assert isinstance(non_blocking, bool) + + return TensorMeta(a, device=utils.canonicalize_device(device)) + + +def _device_put_aten( + a: Tensor, device: Union[str, torch.device], non_blocking=False +) -> Tensor: + return a.to(device, non_blocking=non_blocking) + + +_device_put_doc = """ + Creates a copy of a tensor on the given device. + """ + +device_put = _make_prim( + schema="device_put(Tensor a, Device device, bool non_blocking=False) -> Tensor", + meta=_device_put_meta, + impl_aten=_device_put_aten, + return_type=RETURN_TYPE.NEW, + doc=_device_put_doc, +) + + +# NOTE: need to model meta scalars +# See https://github.com/pytorch/pytorch/issues/78070 +def _item_meta(a: TensorLikeType) -> FakeTensor: + number_type = utils.dtype_to_type(a.dtype) + return TensorMeta(number_type(-1)) + + +_item_doc = """ + Converts a tensor with one element to a Python number. +""" + + +# We can't call into python dispatcher for item again +# because the current prim decomp calls into python dispatcher +# again. https://github.com/pytorch/pytorch/issues/136050 +def _item_aten_no_python_dispatcher(*args, **kwargs): + with torch._dispatch.python.no_python_dispatcher(): + return torch.Tensor.item(*args, **kwargs) + + +# TODO: create a new return type for scalars? +# FIXME: currently returns integers for boolean tensors +# https://github.com/pytorch/pytorch/issues/78071 +item = _make_prim( + schema="item(Tensor a) -> Scalar", + meta=_item_meta, + impl_aten=_item_aten_no_python_dispatcher, + return_type=RETURN_TYPE.NEW, + doc=_item_doc, +) + + +# NOTE: need to model meta scalars +# See https://github.com/pytorch/pytorch/issues/78070 +def _maximum_value_meta(dtype: torch.dtype) -> FakeTensor: + number_type = utils.dtype_to_type(dtype) + return TensorMeta(number_type(-1)) + + +def _maximum_value_aten(dtype: torch.dtype): + if dtype == torch.bool: + return True + elif dtype.is_complex or dtype.is_floating_point: + return torch.finfo(dtype).max + else: + return torch.iinfo(dtype).max + + +_maximum_value_doc = """ + Return the maximum finite value for a dtype. +""" + +# TODO: create a new return type for scalars? +# FIXME: currently returns integers for boolean tensors +# https://github.com/pytorch/pytorch/issues/78071 +maximum_value = _make_prim( + schema="maximum_value(ScalarType dtype) -> Scalar", + meta=_maximum_value_meta, + impl_aten=_maximum_value_aten, + return_type=RETURN_TYPE.NEW, + doc=_maximum_value_doc, +) + + +# NOTE: need to model meta scalars +# See https://github.com/pytorch/pytorch/issues/78070 +def _minimum_value_meta(dtype: torch.dtype) -> FakeTensor: + number_type = utils.dtype_to_type(dtype) + return TensorMeta(number_type(-1)) + + +def _minimum_value_aten(dtype: torch.dtype): + if dtype == torch.bool: + return False + elif dtype.is_complex or dtype.is_floating_point: + return torch.finfo(dtype).min + else: + return torch.iinfo(dtype).min + + +_minimum_value_doc = """ + Return the minimum finite value for a dtype. +""" + +# TODO: create a new return type for scalars? +# FIXME: currently returns integers for boolean tensors +# https://github.com/pytorch/pytorch/issues/78071 +minimum_value = _make_prim( + schema="minimum_value(ScalarType dtype) -> Scalar", + meta=_minimum_value_meta, + impl_aten=_minimum_value_aten, + return_type=RETURN_TYPE.NEW, + doc=_minimum_value_doc, +) + +# +# Inplace operators +# + + +def _copy_to_meta(a: TensorLikeType, b: TensorLikeType): + assert isinstance(a, TensorLike) + assert isinstance(b, TensorLike) + + # Validates the cast is safe + # TODO: move this as an option on the reference + # a_typ = utils.dtype_to_type(a.dtype) + # b_typ = utils.dtype_to_type(b.dtype) + # if a_typ is not utils.get_higher_type(a_typ, b_typ): + # raise RuntimeError(str(b.dtype), " can't be cast safely to ", str(a.dtype), "!") + + # Validates the tensors have the same number of elements + if a.numel() != b.numel(): + msg = f"Attempting to copy {b.numel()} elements to a tensor with {a.numel()} elements!" + raise RuntimeError(msg) + + return a + + +def _copy_to_aten(a: Tensor, b: Tensor) -> Tensor: + return a.copy_(b) + + +_copy_to_doc = """ + Copies the data in b to a and returns the modified a. + """ + +# TODO: Remove safe casting and implement on reference instead +copy_to = _make_prim( + schema="copy_to(Tensor(a!) a, Tensor b) -> Tensor(a!)", + meta=_copy_to_meta, + impl_aten=_copy_to_aten, + return_type=RETURN_TYPE.INPLACE, + doc=_copy_to_doc, + register_conj_neg_fallthrough=True, +) + + +def _copy_strided_meta(a: TensorLikeType, stride: ShapeType): + assert isinstance(a, TensorLike) + return torch.empty_strided( + a.shape, + stride, + dtype=a.dtype, + layout=a.layout, + device=a.device, + requires_grad=a.requires_grad, + ) + + +def _copy_strided_aten(a: Tensor, stride: ShapeType) -> Tensor: + out = torch.empty_strided( + a.size(), + stride=stride, + dtype=a.dtype, + layout=a.layout, + device=a.device, + requires_grad=a.requires_grad, + ) + out.copy_(a) + return out + + +_copy_strided_doc = """ + Copies the data in a to a new tensor, the new tensor has same shape with a size, but has different stride. + """ + + +copy_strided = _make_prim( + schema="copy_strided(Tensor a, SymInt[] stride) -> Tensor", + meta=_copy_strided_meta, + impl_aten=_copy_strided_aten, + return_type=RETURN_TYPE.NEW, + doc=_copy_strided_doc, +) + + +def _resize_meta(a: TensorLikeType, shape: ShapeType): + return a.resize_(shape) + + +def _resize_aten(a: Tensor, shape: ShapeType) -> Tensor: + return a.resize_(shape) + + +_resize_doc = """ + Gives a tensor with no elements a new shape, returning the modified tensor. + + The tensor's strides are contiguous and its values are unitialized. + """ + +# TODO: review support arbitrary resizes +resize = _make_prim( + schema="resize(Tensor(a!) a, SymInt[] shape) -> Tensor(a!)", + meta=_resize_meta, + impl_aten=_resize_aten, + return_type=RETURN_TYPE.INPLACE, + doc=_resize_doc, +) + + +def _reduction_meta(inp, dims, *, output_dtype=None): + """ + Meta function for single output reduction operations + Stride logic is incorrect + """ + assert isinstance(inp, TensorLike) + if output_dtype is None: + output_dtype = inp.dtype + output_shape = utils.compute_reduction_output_shape(inp.shape, dims) + return TensorMeta( + shape=output_shape, + strides=utils.make_contiguous_strides_for(output_shape), + dtype=output_dtype, + device=inp.device, + ) + + +def _var_reduction_meta(inp, dims, correction): + if utils.is_complex_dtype(inp.dtype): + output_dtype = utils.corresponding_real_dtype(inp.dtype) + else: + output_dtype = inp.dtype + return _reduction_meta(inp, dims, output_dtype=output_dtype) + + +_sum_doc = """ + Computes the sum of elements in the input tensor over the list of dimensions + specified in the dim argument + """ +_xor_sum_doc = """ + Computes the xor sum of elements in the input tensor over the list of dimensions + specified in the dim argument + """ +_prod_doc = """ + Computes the product of elements in the input tensor over the list of dimensions + specified in the dim argument + """ +_amax_doc = """ + Computes the maximum value of elements in the input tensor over the list of dimensions + specified in the dim argument + """ +_amin_doc = """ + Computes the minimum value of elements in the input tensor over the list of dimensions + specified in the dim argument + """ +_var_doc = """ + Computes the biased variance of x over the list of dimensions specified in the dim argument + """ + + +def _make_reduction_prim(name: str, impl_aten, doc): + """Creates a reduction prim.""" + return _make_prim( + schema=f"{name}(Tensor inp, int[]? dims, *, ScalarType? output_dtype=None) -> Tensor", + meta=_reduction_meta, + impl_aten=impl_aten, + return_type=RETURN_TYPE.NEW, + doc=doc, + ) + + +def _make_var_reduction_prim(name: str, impl_aten, doc): + """Creates a reduction prim.""" + return _make_prim( + schema=f"{name}(Tensor inp, int[]? dims, float? correction=1, *, ScalarType? output_dtype=None) -> Tensor", + meta=_var_reduction_meta, + impl_aten=impl_aten, + return_type=RETURN_TYPE.NEW, + doc=doc, + ) + + +sum = _make_reduction_prim( + name="sum", + impl_aten=torch.sum, + doc=_sum_doc, +) + + +def _xor_sum_aten( + inp: TensorLikeType, + dims: Optional[DimsSequenceType], + *, + dtype: Optional[torch.dtype] = None, +) -> Tensor: + raise NotImplementedError("xor_sum only implemented with inductor") + + +xor_sum = _make_reduction_prim( + name="xor_sum", + impl_aten=_xor_sum_aten, + doc=_xor_sum_doc, +) + + +def _prod_aten( + inp: TensorLikeType, + dims: Optional[DimsSequenceType], + *, + dtype: Optional[torch.dtype] = None, +) -> Tensor: + if dims is not None: + if len(dims) == 0: + return inp.clone() + for d in sorted(dims, reverse=True): + assert d >= 0 + inp = torch.prod(inp, d, dtype=dtype) + return inp + else: + return torch.prod(inp, dims, dtype=dtype) + + +prod = _make_reduction_prim( + name="prod", + impl_aten=_prod_aten, + doc=_prod_doc, +) + + +# torch.var, but correction is not kwarg-only +def torch_var(input, dim=None, correction=1, **kwargs): + return torch.var(input, dim=dim, correction=correction, **kwargs) + + +var = _make_var_reduction_prim( + name="var", + impl_aten=torch_var, + doc=_var_doc, +) + +amax = _make_reduction_prim( + name="amax", + impl_aten=torch.amax, + doc=_amax_doc, +) + +amin = _make_reduction_prim( + name="amin", + impl_aten=torch.amin, + doc=_amin_doc, +) + + +_iota_doc = """ + Constructs a 1-D tensor t where ``t[i] == start + i * step``. +""" + + +# TODO: layout, pin_memory, memory_format +# TODO: model requires_grad on TensorMeta +def _iota_meta( + length: int, + *, + start: int, + step: int, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, +) -> TensorLikeType: + torch._check( + utils.is_integer_dtype(dtype), + lambda: "prims.iota only supports integer dtypes", + ) + torch._check(step != 0, lambda: "step must be nonzero") + return torch.empty( + length, + dtype=dtype, + device=device, + requires_grad=requires_grad, + ) + + +def _iota_aten( + length: int, + *, + start: int, + step: int, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, +) -> TensorLikeType: + end = start + length * step + return torch.arange( + start, end, step, dtype=dtype, device=device, requires_grad=requires_grad + ) + + +iota = _make_prim( + schema="iota(SymInt length, *, SymInt start, SymInt step, ScalarType dtype, Device device, bool requires_grad) -> Tensor", # noqa: B950 + return_type=RETURN_TYPE.NEW, + meta=_iota_meta, + impl_aten=_iota_aten, + doc=_iota_doc, +) + + +# TODO: layout, pin_memory, memory_format +# TODO: model requires_grad on TensorMeta +def _empty_meta( + shape: ShapeType, *, dtype: torch.dtype, device: torch.device, requires_grad: bool +) -> TensorLikeType: + strides = utils.make_contiguous_strides_for(shape) + return TensorMeta(shape=shape, strides=strides, dtype=dtype, device=device) + + +def _empty_aten( + shape: ShapeType, *, dtype: torch.dtype, device: torch.device, requires_grad: bool +) -> Tensor: + return torch.empty(shape, dtype=dtype, device=device, requires_grad=requires_grad) + + +_empty_doc = """ + Creates a tensor with uninitialized values and the specified shape, dtype, and device. +""" + +empty = _make_prim( + schema="empty(SymInt[] shape, *, ScalarType dtype, Device device, bool requires_grad) -> Tensor", + meta=_empty_meta, + impl_aten=_empty_aten, + return_type=RETURN_TYPE.NEW, + doc=_empty_doc, +) + + +def _empty_strided_meta( + shape: ShapeType, + strides: StrideType, + *, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, +) -> TensorLikeType: + return TensorMeta(shape=shape, strides=strides, dtype=dtype, device=device) + + +_empty_strided_doc = """ + Creates a tensor with uninitialized values. +""" + +# TODO: add layout, pin_memory +empty_strided = _make_prim( + schema="empty_strided(SymInt[] shape, SymInt[] strides, *, ScalarType dtype, Device device, bool requires_grad) -> Tensor", + return_type=RETURN_TYPE.NEW, + meta=_empty_strided_meta, + impl_aten=torch.empty_strided, + doc=_empty_strided_doc, +) + + +def _empty_permuted_meta( + shape: ShapeType, + physical_layout: DimsSequenceType, + *, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, +) -> TensorLikeType: + p_strides = utils.make_contiguous_strides_for([shape[l] for l in physical_layout]) + dim = len(shape) + torch._check( + len(physical_layout) == dim, + lambda: ( + "Number of dimensions in the tensor input does not match the " + f"length of the physical layout; i.e. len(size) = {dim} " + f"is not equal to len(physical_layout) = {len(physical_layout)}" + ), + ) + strides = [0] * len(shape) + seen_dims = set() + for p, l in enumerate(physical_layout): + torch._check( + 0 <= l < dim, + lambda: ( + f"Dimension out of range (expected to be between 0 and {dim - 1}, but got " + f"{l} at index {p}). NB: negative dims " + "not currently supported; file an issue if you want it." + ), + ) + torch._check(l not in seen_dims, lambda: "Duplicate dim not allowed") + strides[l] = p_strides[p] + seen_dims.add(l) + return TensorMeta( + shape=shape, + strides=strides, + dtype=dtype, + device=device, + ) + + +_empty_permuted_doc = """ + Creates a tensor with uninitialized values according to some physical layout, + that is guaranteed to be non-overlapping and dense. +""" + +# TODO: add layout, pin_memory +empty_permuted = _make_prim( + schema="empty_permuted(SymInt[] shape, int[] physical_layout, *, ScalarType dtype, Device device, bool requires_grad) -> Tensor", # noqa: B950 + return_type=RETURN_TYPE.NEW, + meta=_empty_permuted_meta, + impl_aten=torch.empty_permuted, + doc=_empty_permuted_doc, +) + + +def _full_meta( + shape: ShapeType, + fill_value: NumberType, + *, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, +) -> TensorLikeType: + strides = utils.make_contiguous_strides_for(shape) + return TensorMeta(shape=shape, strides=strides, dtype=dtype, device=device) + + +def _full_aten( + shape: ShapeType, + fill_value: NumberType, + *, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, +) -> Tensor: + # Note that Mypy thinks torch.full can't accept a complex fill_value + return torch.full( + shape, fill_value, dtype=dtype, device=device, requires_grad=requires_grad # type: ignore[arg-type] + ) + + +_full_doc = """ + Creates a tensor filled with the given fill value, and with the specified shape, dtype, and device. +""" + +# TODO: add layout +full = _make_prim( + schema="full(SymInt[] shape, Scalar fill_value, *, ScalarType dtype, Device device, bool requires_grad) -> Tensor", + meta=_full_meta, + impl_aten=_full_aten, + return_type=RETURN_TYPE.NEW, + doc=_full_doc, +) + + +def _full_like_meta( + a: TensorLikeType, + fill_value: NumberType, + *, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, +) -> TensorLikeType: + strides = utils.compute_elementwise_output_strides(a) + if a.numel() == 0: + strides = a.stride() + + return TensorMeta(a, strides=strides, dtype=dtype, device=device) + + +def _full_like_aten( + a: Tensor, + fill_value: NumberType, + *, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, +) -> Tensor: + # Note that Mypy thinks torch.full can't accept a complex fill_value + return torch.full_like( + a, fill_value, dtype=dtype, device=device, requires_grad=requires_grad # type: ignore[arg-type] + ) + + +_full_like_doc = """ + Creates a tensor filled with the given fill value, and the same shape, dtype, and device as the + given tensor by default. The dtype and device settings can be overridden + by specifying them explicitly. +""" + +full_like = _make_prim( + schema="full_like(Tensor a, Scalar fill_value, *, ScalarType dtype, Device device, bool requires_grad) -> Tensor", + meta=_full_like_meta, + impl_aten=_full_like_aten, + return_type=RETURN_TYPE.NEW, + doc=_full_like_doc, +) + + +def _scalar_tensor_meta( + scalar: NumberType, + *, + dtype: torch.dtype, + device: torch.device, +) -> TensorLikeType: + shape: ShapeType = [] + strides = utils.make_contiguous_strides_for(shape) + return TensorMeta(scalar, shape=shape, strides=strides, dtype=dtype, device=device) + + +def _scalar_tensor_aten( + scalar: NumberType, + *, + dtype: torch.dtype, + device: torch.device, +) -> Tensor: + if isinstance(scalar, complex) and ( + dtype is None or not utils.is_complex_dtype(dtype) + ): + raise TypeError("Complex scalar requires complex tensor dtype.") + # Note that Mypy thinks torch.scalar can't accept a complex scalar + return torch.scalar_tensor(scalar, dtype=dtype, device=device) # type: ignore[arg-type] + + +_scalar_tensor_doc = """ + Wraps a Number into a Tensor with the specified dtype and device. +""" + +# TODO: add layout and pin_memory support +scalar_tensor = _make_prim( + schema="scalar_tensor(Scalar s, *, ScalarType? dtype=None, Device? device=None) -> Tensor", + meta=_scalar_tensor_meta, + impl_aten=_scalar_tensor_aten, + return_type=RETURN_TYPE.NEW, + doc=_scalar_tensor_doc, +) + + +# +# Linear algebra (linalg) prims +# + + +def _svd_meta( + A: TensorLikeType, *, full_matrices: bool +) -> Tuple[TensorLikeType, TensorLikeType, TensorLikeType]: + utils.check_is_matrix(A, "linalg.svd") + utils.check_fp_or_complex(A.dtype, "linalg.svd", allow_low_precision_dtypes=False) + + A_shape = A.shape + batch = A_shape[:-2] + m, n = A_shape[-2:] + k = min(m, n) + + shape_U = batch + (m, m if full_matrices else k) + strides_U = utils.make_contiguous_strides_for(shape_U, row_major=False) + U = TensorMeta(shape=shape_U, strides=strides_U, dtype=A.dtype, device=A.device) + + shape_S = batch + (k,) + strides_S = utils.make_contiguous_strides_for(shape_S) + S = TensorMeta( + shape=shape_S, + strides=strides_S, + dtype=utils.corresponding_real_dtype(A.dtype) if A.is_complex() else A.dtype, + device=A.device, + ) + + shape_Vh = batch + (n if full_matrices else k, n) + # The CPU backend returns V, but the cuSolver backend returns V^H + # TODO The MAGMA backend returns V, so this is wrong if used with the MAGMA backend + is_cuda = A.device.type == "cuda" + strides_Vh = utils.make_contiguous_strides_for(shape_Vh, row_major=is_cuda) + Vh = TensorMeta(shape=shape_Vh, strides=strides_Vh, dtype=A.dtype, device=A.device) + # Also makes sure this is CUDA or HIP: + # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip + if A.numel() != 0 and Vh.is_complex() and torch.cuda.is_available(): + Vh = Vh.conj() + return U, S, Vh + + +def _svd_aten( + A: TensorLikeType, *, full_matrices: bool +) -> Tuple[Tensor, Tensor, Tensor]: + return torch.linalg.svd(A, full_matrices=full_matrices) + + +_svd_doc = """ + Returns the SVD of a matrix or batch of matrices. + + The `full_matrices` flag controls whether the full or reduced SVD decomposition is returned. +""" + +svd = _make_prim( + schema="svd(Tensor A, *, bool full_matrices) -> (Tensor U, Tensor S, Tensor Vh)", + meta=_svd_meta, + impl_aten=_svd_aten, + return_type=(RETURN_TYPE.NEW, RETURN_TYPE.NEW, RETURN_TYPE.NEW), + doc=_svd_doc, +) + + +# +# Randomness Prims +# + + +def _normal_meta( + shape: ShapeType, + *, + mean: Union[float, complex], + std: float, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, + generator: Optional[torch.Generator] = None, +) -> TensorLikeType: + torch._check( + std >= 0.0, + lambda: f"expected non-negative standard deviation, but got std={std}", + ) + + torch._check( + utils.is_float_dtype(dtype) or utils.is_complex_dtype(dtype), + lambda: f"expected a floating-point or complex dtype, but got dtype={dtype}", + ) + + strides = utils.make_contiguous_strides_for(shape) + return TensorMeta(shape=shape, strides=strides, dtype=dtype, device=device) + + +def _normal_aten( + shape: ShapeType, + *, + mean: Union[float, complex], + std: float, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, + generator: Optional[torch.Generator] = None, +) -> Tensor: + a = torch.empty(shape, dtype=dtype, device=device, requires_grad=requires_grad) + with torch.no_grad(): + # NOTE: normal_ is incorrectly annotated to expect mean to be a float + a.normal_(mean, std, generator=generator) # type: ignore[arg-type] + return a + + +_normal_doc = """ + Constructs a tensor filled with values drawn from a normal distribution with the specified mean + and standard deviation. + + Only supports floating-point types. +""" + +normal = _make_prim( + schema=( + "normal(SymInt[] shape, *, Scalar mean, Scalar std, ScalarType dtype, Device device, bool requires_grad, Generator? generator=None) -> Tensor" # noqa: B950 + ), + return_type=RETURN_TYPE.NEW, + meta=_normal_meta, + impl_aten=_normal_aten, + doc=_normal_doc, +) + + +def _uniform_meta( + shape: ShapeType, + *, + low: float, + high: float, + dtype: torch.dtype, + device: torch.device, + generator: Optional[torch.Generator] = None, +) -> TensorLikeType: + strides = utils.make_contiguous_strides_for(shape) + return TensorMeta(shape=shape, strides=strides, dtype=dtype, device=device) + + +def _uniform_aten( + shape: ShapeType, + *, + low: float, + high: float, + dtype: torch.dtype, + device: torch.device, + generator: Optional[torch.Generator] = None, +) -> Tensor: + a = torch.empty(shape, dtype=dtype, device=device) + a.uniform_(low, high, generator=generator) + return a + + +_uniform_doc = """ + Constructs a tensor filled with values drawn uniformly from low to high. +""" + +# TODO: we should more seriously review randomness modeling and prims +_uniform_helper = _make_prim( + schema=( + "uniform(SymInt[] shape, *, Scalar low, Scalar high, ScalarType dtype, Device device, Generator? generator=None) -> Tensor" + ), + return_type=RETURN_TYPE.NEW, + meta=_uniform_meta, + impl_aten=_uniform_aten, + doc=_uniform_doc, +) + +# +# FFT prims +# + + +def _fft_r2c_meta( + input: TensorLike, + *, + dim: DimsSequenceType, + onesided: bool, +) -> TensorLikeType: + dim = utils.canonicalize_dims(input.ndim, dim) + utils.validate_no_repeating_dims(dim) + + shape = list(input.shape) + if onesided: + last_dim = dim[-1] + shape[last_dim] = shape[last_dim] // 2 + 1 + + dtype = utils.corresponding_complex_dtype(input.dtype) + strides = utils.make_contiguous_strides_for(shape) + return TensorMeta(shape=shape, strides=strides, dtype=dtype, device=input.device) + + +def _fft_r2c_aten( + input: TensorLike, + *, + dim: DimsSequenceType, + onesided: bool, +) -> TensorLikeType: + normalization = 0 # No normalization + return torch._fft_r2c(input, dim, normalization, onesided) + + +_fft_r2c_doc = """ + Performs a real to complex Fast Fourier Transform +""" + + +fft_r2c = _make_prim( + schema="fft_r2c(Tensor self, *, int[] dim, bool onesided) -> Tensor", + meta=_fft_r2c_meta, + impl_aten=_fft_r2c_aten, + return_type=RETURN_TYPE.NEW, + doc=_fft_r2c_doc, +) + + +def _fft_c2c_meta( + input: TensorLike, + *, + dim: DimsSequenceType, + forward: bool, +) -> TensorLikeType: + dim = utils.canonicalize_dims(input.ndim, dim) + utils.validate_no_repeating_dims(dim) + + shape = input.shape + strides = utils.make_contiguous_strides_for(shape) + return TensorMeta( + shape=shape, strides=strides, dtype=input.dtype, device=input.device + ) + + +def _fft_c2c_aten( + input: TensorLike, + *, + dim: DimsSequenceType, + forward: bool, +) -> TensorLikeType: + normalization = 0 # No normalization + return torch._fft_c2c(input, dim, normalization, forward) + + +_fft_c2c_doc = """ + Performs either a Fast Fourier Transform, or its inverse +""" + + +fft_c2c = _make_prim( + schema="fft_c2c(Tensor self, *, int[] dim, bool forward) -> Tensor", + meta=_fft_c2c_meta, + impl_aten=_fft_c2c_aten, + return_type=RETURN_TYPE.NEW, + doc=_fft_c2c_doc, +) + + +def _fft_c2r_meta( + input: TensorLike, + *, + dim: DimsSequenceType, + last_dim_size: int, +) -> TensorLikeType: + dim = utils.canonicalize_dims(input.ndim, dim) + utils.validate_no_repeating_dims(dim) + + shape = list(input.shape) + shape[dim[-1]] = last_dim_size + dtype = utils.corresponding_real_dtype(input.dtype) + strides = utils.make_contiguous_strides_for(shape) + return TensorMeta(shape=shape, strides=strides, dtype=dtype, device=input.device) + + +def _fft_c2r_aten( + input: TensorLike, + *, + dim: DimsSequenceType, + last_dim_size: int, +) -> TensorLikeType: + normalization = 0 # No normalization + return torch._fft_c2r(input, dim, normalization, last_dim_size) + + +_fft_c2r_doc = """ + Performs a complex to real Inverse Fast Fourier Transform +""" + + +fft_c2r = _make_prim( + schema="fft_c2r(Tensor self, *, int[] dim, SymInt last_dim_size) -> Tensor", + meta=_fft_c2r_meta, + impl_aten=_fft_c2r_aten, + return_type=RETURN_TYPE.NEW, + doc=_fft_c2r_doc, +) + + +def _frexp_meta(self: TensorLikeType) -> Tuple[TensorLikeType, TensorLikeType]: + torch._check( + self.dtype.is_floating_point, + lambda: "torch.frexp() only supports floating-point dtypes", + ) + return torch.empty_like(self), torch.empty_like(self, dtype=torch.int32) + + +frexp = _make_prim( + schema="frexp(Tensor self) -> (Tensor mantissa, Tensor exponent)", + meta=_frexp_meta, + return_type=(RETURN_TYPE.NEW, RETURN_TYPE.NEW), + impl_aten=torch.frexp, + doc="", +) + + +def _make_token_aten() -> TensorLikeType: + return new_token_tensor() + + +_make_token = _make_prim( + schema="_make_token() -> Tensor", + meta=_make_token_aten, + return_type=RETURN_TYPE.NEW, + impl_aten=_make_token_aten, + doc="Creates a token used for keeping track of side effects.", +) + + +def _sink_tokens_aten(tokens) -> None: + pass + + +_sink_tokens = _make_prim( + schema="_sink_tokens(Tensor[] tokens) -> ()", + meta=_sink_tokens_aten, + return_type=RETURN_TYPE.NONE, + impl_aten=_sink_tokens_aten, + doc="Sink all of the tokens which were previously used for keeping track of side effects.", +) + + +register_rng_prims() +register_debug_prims() diff --git a/lib/python3.12/site-packages/torch/_prims/context.py b/lib/python3.12/site-packages/torch/_prims/context.py new file mode 100644 index 0000000000000000000000000000000000000000..7026d02790715b01bfd33a7190062d4c83b7ba4f --- /dev/null +++ b/lib/python3.12/site-packages/torch/_prims/context.py @@ -0,0 +1,146 @@ +# mypy: allow-untyped-defs +import functools +from contextlib import nullcontext +from typing import Any, Callable, Dict, Optional, Sequence + +import torch +import torch._decomp +import torch._prims +import torch._refs +import torch._refs.nn +import torch._refs.nn.functional +import torch._refs.special +import torch.overrides +from torch._prims_common import torch_function_passthrough + + +@functools.lru_cache(None) +def torch_to_refs_map(): + """ + Mapping of torch API functions to torch._refs functions. + E.g. torch_to_refs_map()[torch.add] == torch._refs.add + """ + modules = [ + (torch, torch._refs), + (torch.nn, torch._refs.nn), + (torch.nn.functional, torch._refs.nn.functional), + (torch.special, torch._refs.special), + (torch.fft, torch._refs.fft), + (torch.linalg, torch._refs.linalg), + ] + r: Dict[Any, Any] = { + torch.Tensor.__invert__: torch._refs.bitwise_not, + torch.Tensor.__xor__: torch._refs.bitwise_xor, + torch.Tensor.__and__: torch._refs.bitwise_and, + torch.Tensor.__or__: torch._refs.bitwise_or, + torch.Tensor.__eq__: torch._refs.eq, + torch.Tensor.__rsub__: torch._refs.rsub, + torch.Tensor.__rtruediv__: torch._refs.rtruediv, + torch.Tensor.__floordiv__: torch._refs.floor_divide, + torch.Tensor.__rfloordiv__: torch._refs.rfloordiv, + torch.Tensor.__pow__: torch._refs.pow, + torch.Tensor.__rpow__: torch._refs.rpow, + torch.Tensor.new_empty: torch._refs.new_empty, + torch.Tensor.new_full: torch._refs.new_full, + torch.Tensor.new_zeros: torch._refs.new_zeros, + torch.Tensor.new_ones: torch._refs.new_ones, + torch.Tensor.fill_: torch._refs.fill_, + torch.Tensor.zero_: torch._refs.zero_, + torch.Tensor.to: torch._refs.to, + torch.Tensor.sum_to_size: torch._refs.sum_to_size, + # TODO: Should these methods be mapped some other way? + torch.Tensor.copy_: torch._prims.copy_to, + torch.Tensor.resize: torch._prims.resize, + } + for mod_torch, mod_refs in modules: + for s in mod_refs.__all__: # type: ignore[attr-defined] + r[mod_torch.__dict__.get(s)] = mod_refs.__dict__.get(s) + + # Support remapping torch.Tensor.foo to _refs.foo + for s in dir(torch.Tensor): + if s in torch._refs.__all__: + r[getattr(torch.Tensor, s)] = torch._refs.__dict__.get(s) + + # Support conversions + for s in torch._refs._conversions.__all__: + tensor_attr = getattr(torch.Tensor, s, None) or getattr(torch, s) + r[tensor_attr] = torch._refs._conversions.__dict__.get(s) + + return r + + +@functools.lru_cache(None) +def all_prims(): + """ + Set of all prim functions, e.g., torch._prims.add in all_prims() + """ + return {torch._prims.__dict__.get(s) for s in torch._prims.__all__} + + +class TorchRefsMode(torch.overrides.TorchFunctionMode): + """ + Switches the interpretation of torch.* functions and Tensor methods to + use PrimTorch refs in torch._refs. (Direct calls to _refs are unaffected.) + + >>> # xdoctest: +SKIP + >>> with TorchRefsMode(): + ... torch.add(x, y) # calls torch._refs.add(x, y) + + By default, this context manager will fall back on the torch.* if the + ref does not exist; set strict=True to error if this occurs. + If the ref exists we still would like to fall back on the torch.* sometimes, + this behavior can be customized by passing a function to should_fallback_fn. + """ + + def __init__( + self, + strict=False, + should_fallback_fn=lambda *_: False, + prims_mode_cls=nullcontext, + ): + self.strict = strict + self.should_fallback_fn = should_fallback_fn + self.prims_mode_cls = prims_mode_cls + + def __torch_function__( + self, + orig_func: Callable, + types: Sequence, + args: Sequence[Any] = (), + kwargs: Optional[Dict] = None, + ): + if kwargs is None: + kwargs = {} + # For primitive operations, run them as is without interception + # Unless we are in prims_mode, in which case we want to use nvprims + if orig_func in torch_function_passthrough or orig_func in all_prims(): + with self.prims_mode_cls(): + return orig_func(*args, **kwargs) + mapping = torch_to_refs_map() + func = mapping.get(orig_func, None) + + # For torch.ops.aten.*, use registered decompositions from torch._decomp + # torch._decomp.decomposition_table provides a mapping from + # torch.ops.aten.* to torch._refs or torch._decomp.decompositions + # implementations. + # There're other ways to implement this functionality, + # see https://github.com/pytorch/pytorch/pull/82657#discussion_r939776417 + if func is None and isinstance(orig_func, torch._ops.OpOverload): + func = torch._decomp.decomposition_table.get(orig_func, None) + elif func is None and isinstance(orig_func, torch._ops.OpOverloadPacket): + default = getattr(orig_func, "default", None) + if default is not None: + func = torch._decomp.decomposition_table.get(default, None) + + if func is not None: + # If the ref exists query whether we should use it or not + if self.should_fallback_fn(self, orig_func, func, args, kwargs): + return orig_func(*args, **kwargs) + # torch calls inside func should be interpreted as refs calls + with self: + return func(*args, **kwargs) + if self.strict: + raise RuntimeError( + f"no _refs support for {torch.overrides.resolve_name(orig_func)}" + ) + return orig_func(*args, **kwargs) diff --git a/lib/python3.12/site-packages/torch/_prims/debug_prims.py b/lib/python3.12/site-packages/torch/_prims/debug_prims.py new file mode 100644 index 0000000000000000000000000000000000000000..6958cbcef283dbc7eda544760640126e7524bb2e --- /dev/null +++ b/lib/python3.12/site-packages/torch/_prims/debug_prims.py @@ -0,0 +1,54 @@ +# mypy: allow-untyped-defs +import contextlib +from typing import Optional + +import torch +from torch.utils._content_store import ContentStoreReader + + +LOAD_TENSOR_READER: Optional[ContentStoreReader] = None + + +@contextlib.contextmanager +def load_tensor_reader(loc): + global LOAD_TENSOR_READER + assert LOAD_TENSOR_READER is None + # load_tensor is an "op", and we will play merry hell on + # Inductor's memory planning if we return a tensor that + # aliases another tensor that we previously returned from + # an operator. So unlike standard ContentStoreReader use, + # we disable the cache so that you always get fresh storages + # (no aliasing for you!) + LOAD_TENSOR_READER = ContentStoreReader(loc, cache=False) + try: + yield + finally: + LOAD_TENSOR_READER = None + + +def register_debug_prims(): + torch.library.define( + "debugprims::load_tensor", + "(str name, int[] size, int[] stride, *, ScalarType dtype, Device device) -> Tensor", + ) + + @torch.library.impl("debugprims::load_tensor", "BackendSelect") + def load_tensor_factory(name, size, stride, dtype, device): + if LOAD_TENSOR_READER is None: + from torch._dynamo.testing import rand_strided + + return rand_strided(size, stride, dtype, device) + else: + from torch._dynamo.utils import clone_input + + # device argument here takes care of coercion + r = LOAD_TENSOR_READER.read_tensor(name, device=device) + assert list(r.size()) == size, f"{r.size()} != {size}" + assert list(r.stride()) == stride, f"{r.stride()} != {stride}" + assert r.device == device, f"{r.device} != {device}" + + # Unlike the other properties, we will do coercions for dtype + # mismatch + if r.dtype != dtype: + r = clone_input(r, dtype=dtype) + return r diff --git a/lib/python3.12/site-packages/torch/_prims/executor.py b/lib/python3.12/site-packages/torch/_prims/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..a020c1db75b2eacb4cb4c6ca809586258e74dcd8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_prims/executor.py @@ -0,0 +1,60 @@ +# mypy: allow-untyped-defs +from typing import Callable, Optional + +from torch._prims.context import TorchRefsMode +from torch.fx import GraphModule +from torch.fx.experimental.proxy_tensor import make_fx, wrapper_and_args_for_make_fx + + +def execute( + gm: GraphModule, + *args, + executor: str = "aten", + executor_parameters: Optional[dict] = None, +): + """ + Prototype ATen executor. + + Just executes the context's graph. + """ + + if executor == "aten": + return gm.forward(*args) + + msg = f"Received unexpected value for 'executor': {executor}. Allowed values are: aten." + raise ValueError(msg) + + +def make_traced(fn: Callable): + """ + Returns a function that, when called, will + trace its torch operations to prims and then + execute those prims on the requested trace executor + (possibly lowering them to that trace executor first). + + Only supports the torch operations defined in _torch_to_reference_map + in context.py and operations with positional args. All args must + be tensors. + In the near future all these restrictions will be lifted. + + Example usage: + + def foo(a, b): + return torch.add(a, b) + + traced_foo = make_traced(foo) + + a = torch.randn((1, 2, 3, 4, 5), device='cuda') + b = torch.randn((1, 2, 3, 4, 5), device='cuda') + result = traced_foo(a, b, executor='aten') + """ + + def _traced(*args, executor="aten", **kwargs): + # TODO: caching + wrapped, all_args = wrapper_and_args_for_make_fx(fn, args, kwargs) + + with TorchRefsMode(): + gm = make_fx(wrapped)(all_args) + return execute(gm, all_args, executor=executor) + + return _traced diff --git a/lib/python3.12/site-packages/torch/_prims/rng_prims.py b/lib/python3.12/site-packages/torch/_prims/rng_prims.py new file mode 100644 index 0000000000000000000000000000000000000000..bbbdb8958f9adb7c215a404f28db8272c0b563e2 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_prims/rng_prims.py @@ -0,0 +1,319 @@ +# mypy: allow-untyped-defs +from typing import Optional, Tuple + +import torch +import torch.utils._pytree as pytree +from torch import _prims +from torch._C import DispatchKey +from torch._higher_order_ops.utils import autograd_not_implemented +from torch._ops import HigherOrderOperator +from torch._prims_common import CUDARngStateHelper, make_contiguous_strides_for +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ProxyTorchDispatchMode, + track_tensor_tree, +) +from torch.types import _device, _dtype + + +def throw_on_non_cuda(device): + raise RuntimeError( + f"You are trying to functionalize a {device.type} RNG operator but {device.type} does not " + f"use Philox/counter-based RNG. Therefore, functionalizing a {device.type} RNG operator is " + "not supported. We are discussing the possibility of a Philox-based RNG implementation for CPU." + ) + + +def register_rng_prim(name, schema, impl_aten, impl_meta, doc, tags=None): + rngprim_def = torch.library.custom_op( + "rngprims::" + name, impl_aten, mutates_args=(), schema=schema + ) + rngprim_def.register_fake(impl_meta) + + prim_packet = getattr(torch._ops.ops.rngprims, name) + prim = prim_packet.default + if tags: + prim._tags = tags + + for p in (prim_packet, prim): + p.__doc__ = doc + p.return_type = torch._prims_common.RETURN_TYPE.NEW # type: ignore[attr-defined] + + p.schema = name + schema + p.impl_aten = impl_aten + p.prim_meta_impl = impl_meta + + +# Philox rand offsets could be shared in future with other philox ops, so +# keeping these functions in global scope. +def philox_rand_offset_meta( + shape: torch.Size, +): + return _prims.TensorLike(torch.tensor(0, dtype=torch.int64)) + + +def philox_rand_offset( + shape: torch.Size, +): + # For impl, look at the function calc_execution_policy in the file + # aten/src/ATen/native/cuda/DistributionTemplates.h. The impl was copied at + # commit hash 72aa0667bd16707d50eb8fa337092a1f5d11dfb6 + numel_scalar = 1 + for dim_size in shape: + numel_scalar *= dim_size + numel = torch.scalar_tensor(numel_scalar, dtype=torch.int64) + + block_size = 256 + unroll = 4 + curand4_engine_calls = 4 + device_property = torch.cuda.get_device_properties(torch.cuda.current_device()) + blocks_per_sm = device_property.max_threads_per_multi_processor // block_size + grid_size = (numel + block_size - 1) // block_size + grid_size = min(grid_size, device_property.multi_processor_count * blocks_per_sm) + offset = ( + (numel - 1) // (block_size * grid_size * unroll) + 1 + ) * curand4_engine_calls + return offset + + +def register_philox_rand(): + name = "philox_rand" + schema = "(SymInt[] size, Tensor seed, Tensor offset, int[]? stride, Device? device=None, ScalarType? dtype=None) -> (Tensor, Tensor)" # noqa: B950 + + def _philox_rand_meta( + shape: torch.Size, + seed: torch.Tensor, + offset: torch.Tensor, + stride: Optional[Tuple[int, ...]], + device: _device, + dtype: _dtype, + ): + # stride arg will be useful for distributed usecase. Currently, its unused. + assert stride is None + stride = make_contiguous_strides_for(shape) + random_values = _prims.TensorMeta( + shape=shape, strides=stride, dtype=dtype, device=device + ) + offset = philox_rand_offset_meta(shape) + return (random_values, offset) + + def _philox_rand( + shape: torch.Size, + seed: torch.Tensor, + offset: torch.Tensor, + stride: Optional[Tuple[int, ...]], + device: _device, + dtype: _dtype, + ): + # stride arg will be useful for distributed usecase. Currently, its unused. + assert stride is None + if device.type == "cpu": + devices = [] + else: + devices = [device] + + if device.type != "cuda": + raise throw_on_non_cuda(device) + + with torch.random.fork_rng(devices): + CUDARngStateHelper.set_torch_state_tensor(seed, offset) + random_values = torch.rand(shape, device=device, dtype=dtype) + + return random_values, philox_rand_offset(shape) + + register_rng_prim( + name=name, + schema=schema, + impl_aten=_philox_rand, + impl_meta=_philox_rand_meta, + doc="Philox based stateless rand operator", + tags=(torch.Tag.nondeterministic_seeded,), + ) + + +def get_device(args, kwargs): + if kwargs.get("device"): + device = kwargs.get("device") + if isinstance(device, str): + device = torch.device(device) + return device.type + + devices = {arg.device.type for arg in args if isinstance(arg, torch.Tensor)} + if any(dev == "cuda" for dev in devices): + return "cuda" + elif any(dev == "xpu" for dev in devices): + return "xpu" + elif any(dev == "hpu" for dev in devices): + return "hpu" + elif any(dev == "cpu" for dev in devices): + return "cpu" + return None + + +def register_run_and_save_rng_state_op(): + class RunAndSaveRngState(HigherOrderOperator): + def __init__(self): + super().__init__("run_and_save_rng_state") + + def __call__(self, op, *args, **kwargs): + return super().__call__(op, *args, **kwargs) + + run_and_save_rng_state = RunAndSaveRngState() + + run_and_save_rng_state.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(run_and_save_rng_state, deferred_error=True) + ) + + @run_and_save_rng_state.py_impl(DispatchKey.CUDA) + def impl_cuda(op, *args, **kwargs): + return torch.cuda.get_rng_state(), op(*args, **kwargs) + + @run_and_save_rng_state.py_impl(DispatchKey.CPU) + def impl_cpu(op, *args, **kwargs): + return torch.get_rng_state(), op(*args, **kwargs) + + @run_and_save_rng_state.py_impl(DispatchKey.HPU) + def impl_hpu(op, *args, **kwargs): + if hasattr(torch, "hpu"): + return torch.hpu.get_rng_state(), op(*args, **kwargs) + raise RuntimeError("functionalize a hpu RNG operator is not supported.") + + @run_and_save_rng_state.py_impl(DispatchKey.XPU) + def impl_xpu(op, *args, **kwargs): + return torch.xpu.get_rng_state(), op(*args, **kwargs) + + @run_and_save_rng_state.py_impl(DispatchKey.BackendSelect) + def impl_backend_select(op, *args, **kwargs): + impl_map = { + "cuda": impl_cuda, + "cpu": impl_cpu, + "hpu": impl_hpu, + "xpu": impl_xpu, + } + device = get_device(args, kwargs) + assert device in impl_map, f"Backend not supported for {device}" + impl = impl_map[device] + return impl(op, *args, **kwargs) + + @run_and_save_rng_state.py_impl(FakeTensorMode) + def impl_fake_tensor_mode(mode, op, *args, **kwargs): + # Check device to call the right impl + with mode: + return impl_backend_select(op, *args, **kwargs) + + @run_and_save_rng_state.py_impl(ProxyTorchDispatchMode) + def impl_proxy_dispatch_mode(mode, op, *args, **kwargs): + out = impl_backend_select(op, *args, **kwargs) + proxy_args = pytree.tree_map(mode.tracer.unwrap_proxy, (op, *args)) + proxy_kwargs = pytree.tree_map(mode.tracer.unwrap_proxy, kwargs) + out_proxy = mode.tracer.create_proxy( + "call_function", run_and_save_rng_state, proxy_args, proxy_kwargs + ) + return track_tensor_tree(out, out_proxy, constant=None, tracer=mode.tracer) + + return run_and_save_rng_state + + +def register_run_with_rng_state_op(): + class RunWithRngState(HigherOrderOperator): + def __init__(self): + super().__init__("run_with_rng_state") + + def __call__(self, rng_state, op, *args, **kwargs): + return super().__call__(rng_state, op, *args, **kwargs) + + run_with_rng_state = RunWithRngState() + + run_with_rng_state.py_impl(DispatchKey.Autograd)( + autograd_not_implemented(run_with_rng_state, deferred_error=True) + ) + + @run_with_rng_state.py_impl(DispatchKey.CUDA) + def impl_cuda(rng_state, op, *args, **kwargs): + current_state = torch.cuda.get_rng_state() + torch.cuda.set_rng_state(rng_state.cpu()) + out = op(*args, **kwargs) + torch.cuda.set_rng_state(current_state) + return out + + @run_with_rng_state.py_impl(DispatchKey.CPU) + def impl_cpu(rng_state, op, *args, **kwargs): + current_state = torch.get_rng_state() + torch.set_rng_state(rng_state) + out = op(*args, **kwargs) + torch.set_rng_state(current_state) + return out + + @run_with_rng_state.py_impl(DispatchKey.HPU) + def impl_hpu(rng_state, op, *args, **kwargs): + if hasattr(torch, "hpu"): + current_state = torch.hpu.get_rng_state() + torch.hpu.set_rng_state(rng_state) + out = op(*args, **kwargs) + torch.hpu.set_rng_state(current_state) + return out + raise RuntimeError("functionalize a hpu RNG operator is not supported.") + + @run_with_rng_state.py_impl(DispatchKey.XPU) + def impl_xpu(rng_state, op, *args, **kwargs): + current_state = torch.xpu.get_rng_state() + torch.xpu.set_rng_state(rng_state) + out = op(*args, **kwargs) + torch.xpu.set_rng_state(current_state) + return out + + @run_with_rng_state.py_impl(ProxyTorchDispatchMode) + def impl_proxy_dispatch_mode(mode, rng_state, op, *args, **kwargs): + # TODO: you don't need to do this, the dispatch here already disabled + # it + with disable_proxy_modes_tracing(): + out = run_with_rng_state(rng_state, op, *args, **kwargs) + proxy_args = pytree.tree_map(mode.tracer.unwrap_proxy, (rng_state, op, *args)) + proxy_kwargs = pytree.tree_map(mode.tracer.unwrap_proxy, kwargs) + out_proxy = mode.tracer.create_proxy( + "call_function", run_with_rng_state, proxy_args, proxy_kwargs + ) + return track_tensor_tree(out, out_proxy, constant=None, tracer=mode.tracer) + + @run_with_rng_state.py_impl(DispatchKey.BackendSelect) + def impl_backend_select(rng_state, op, *args, **kwargs): + impl_map = { + "cuda": impl_cuda, + "cpu": impl_cpu, + "hpu": impl_hpu, + "xpu": impl_xpu, + } + device = get_device(args, kwargs) + assert device in impl_map, f"Backend not supported for {device}" + impl = impl_map[device] + return impl(rng_state, op, *args, **kwargs) + + @run_with_rng_state.py_impl(FakeTensorMode) + def impl_fake_tensor_mode(mode, rng_state, op, *args, **kwargs): + # Skip setting the set_rng_state as it does not work well with fake tensors. + # And it does not matter for the fake tensor mode. + with mode: + return op(*args, **kwargs) + + @run_with_rng_state.py_functionalize_impl + def impl_functional(ctx, rng_state, op, *args, **kwargs): + unwrapped_rng_state = ctx.unwrap_tensors(rng_state) + unwrapped_args = ctx.unwrap_tensors(args) + unwrapped_kwargs = ctx.unwrap_tensors(kwargs) + + with ctx.redispatch_to_next(): + out = run_with_rng_state( + unwrapped_rng_state, op, *unwrapped_args, **unwrapped_kwargs + ) + return ctx.wrap_tensors(out) + + return run_with_rng_state + + +run_and_save_rng_state = register_run_and_save_rng_state_op() +run_with_rng_state = register_run_with_rng_state_op() + + +def register_rng_prims(): + register_philox_rand() diff --git a/lib/python3.12/site-packages/torch/backends/__init__.py b/lib/python3.12/site-packages/torch/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3017351888698ec23062348212b0d1884b4f5734 --- /dev/null +++ b/lib/python3.12/site-packages/torch/backends/__init__.py @@ -0,0 +1,72 @@ +# mypy: allow-untyped-defs +import types +from contextlib import contextmanager + + +# The idea for this parameter is that we forbid bare assignment +# to torch.backends..enabled and friends when running our +# test suite, where it's very easy to forget to undo the change +# later. +__allow_nonbracketed_mutation_flag = True + + +def disable_global_flags(): + global __allow_nonbracketed_mutation_flag + __allow_nonbracketed_mutation_flag = False + + +def flags_frozen(): + return not __allow_nonbracketed_mutation_flag + + +@contextmanager +def __allow_nonbracketed_mutation(): + global __allow_nonbracketed_mutation_flag + old = __allow_nonbracketed_mutation_flag + __allow_nonbracketed_mutation_flag = True + try: + yield + finally: + __allow_nonbracketed_mutation_flag = old + + +class ContextProp: + def __init__(self, getter, setter): + self.getter = getter + self.setter = setter + + def __get__(self, obj, objtype): + return self.getter() + + def __set__(self, obj, val): + if not flags_frozen(): + self.setter(val) + else: + raise RuntimeError( + f"not allowed to set {obj.__name__} flags " + "after disable_global_flags; please use flags() context manager instead" + ) + + +class PropModule(types.ModuleType): + def __init__(self, m, name): + super().__init__(name) + self.m = m + + def __getattr__(self, attr): + return self.m.__getattribute__(attr) + + +from torch.backends import ( + cpu as cpu, + cuda as cuda, + cudnn as cudnn, + cusparselt as cusparselt, + mha as mha, + mkl as mkl, + mkldnn as mkldnn, + mps as mps, + nnpack as nnpack, + openmp as openmp, + quantized as quantized, +) diff --git a/lib/python3.12/site-packages/torch/backends/cuda/__init__.py b/lib/python3.12/site-packages/torch/backends/cuda/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2b7aa44946671688aa0eadbeca44a57e2dd4c0e7 --- /dev/null +++ b/lib/python3.12/site-packages/torch/backends/cuda/__init__.py @@ -0,0 +1,480 @@ +# mypy: allow-untyped-defs +import contextlib +from typing import Union +from typing_extensions import deprecated + +import torch + + +__all__ = [ + "is_built", + "cuFFTPlanCacheAttrContextProp", + "cuFFTPlanCache", + "cuFFTPlanCacheManager", + "cuBLASModule", + "preferred_linalg_library", + "preferred_blas_library", + "cufft_plan_cache", + "matmul", + "SDPAParams", + "enable_cudnn_sdp", + "cudnn_sdp_enabled", + "enable_flash_sdp", + "flash_sdp_enabled", + "enable_mem_efficient_sdp", + "mem_efficient_sdp_enabled", + "math_sdp_enabled", + "enable_math_sdp", + "allow_fp16_bf16_reduction_math_sdp", + "fp16_bf16_reduction_math_sdp_allowed", + "is_flash_attention_available", + "can_use_flash_attention", + "can_use_efficient_attention", + "can_use_cudnn_attention", + "sdp_kernel", +] + + +def is_built(): + r""" + Return whether PyTorch is built with CUDA support. + + Note that this doesn't necessarily mean CUDA is available; just that if this PyTorch + binary were run on a machine with working CUDA drivers and devices, we would be able to use it. + """ + return torch._C._has_cuda + + +class cuFFTPlanCacheAttrContextProp: + # Like regular ContextProp, but uses the `.device_index` attribute from the + # calling object as the first argument to the getter and setter. + def __init__(self, getter, setter): + self.getter = getter + self.setter = setter + + def __get__(self, obj, objtype): + return self.getter(obj.device_index) + + def __set__(self, obj, val): + if isinstance(self.setter, str): + raise RuntimeError(self.setter) + self.setter(obj.device_index, val) + + +class cuFFTPlanCache: + r""" + Represent a specific plan cache for a specific `device_index`. + + The attributes `size` and `max_size`, and method `clear`, can fetch and/ or + change properties of the C++ cuFFT plan cache. + """ + + def __init__(self, device_index): + self.device_index = device_index + + size = cuFFTPlanCacheAttrContextProp( + torch._cufft_get_plan_cache_size, + ".size is a read-only property showing the number of plans currently in the " + "cache. To change the cache capacity, set cufft_plan_cache.max_size.", + ) + + max_size = cuFFTPlanCacheAttrContextProp( + torch._cufft_get_plan_cache_max_size, torch._cufft_set_plan_cache_max_size + ) + + def clear(self): + return torch._cufft_clear_plan_cache(self.device_index) + + +class cuFFTPlanCacheManager: + r""" + Represent all cuFFT plan caches, return the cuFFTPlanCache for a given device when indexed. + + Finally, this object, when used directly as a `cuFFTPlanCache` object (e.g., + setting the `.max_size`) attribute, the current device's cuFFT plan cache is + used. + """ + + __initialized = False + + def __init__(self): + self.caches = [] + self.__initialized = True + + def __getitem__(self, device): + index = torch.cuda._utils._get_device_index(device) + if index < 0 or index >= torch.cuda.device_count(): + raise RuntimeError( + f"cufft_plan_cache: expected 0 <= device index < {torch.cuda.device_count()}, but got " + f"device with index {index}" + ) + if len(self.caches) == 0: + self.caches.extend( + cuFFTPlanCache(index) for index in range(torch.cuda.device_count()) + ) + return self.caches[index] + + def __getattr__(self, name): + return getattr(self[torch.cuda.current_device()], name) + + def __setattr__(self, name, value): + if self.__initialized: + return setattr(self[torch.cuda.current_device()], name, value) + else: + return super().__setattr__(name, value) + + +class cuBLASModule: + def __getattr__(self, name): + if name == "allow_tf32": + return torch._C._get_cublas_allow_tf32() + elif name == "allow_fp16_reduced_precision_reduction": + return torch._C._get_cublas_allow_fp16_reduced_precision_reduction() + elif name == "allow_bf16_reduced_precision_reduction": + return torch._C._get_cublas_allow_bf16_reduced_precision_reduction() + raise AttributeError("Unknown attribute " + name) + + def __setattr__(self, name, value): + if name == "allow_tf32": + return torch._C._set_cublas_allow_tf32(value) + elif name == "allow_fp16_reduced_precision_reduction": + return torch._C._set_cublas_allow_fp16_reduced_precision_reduction(value) + elif name == "allow_bf16_reduced_precision_reduction": + return torch._C._set_cublas_allow_bf16_reduced_precision_reduction(value) + raise AttributeError("Unknown attribute " + name) + + +_LinalgBackends = { + "default": torch._C._LinalgBackend.Default, + "cusolver": torch._C._LinalgBackend.Cusolver, + "magma": torch._C._LinalgBackend.Magma, +} +_LinalgBackends_str = ", ".join(_LinalgBackends.keys()) + + +def preferred_linalg_library( + backend: Union[None, str, torch._C._LinalgBackend] = None +) -> torch._C._LinalgBackend: + r""" + Override the heuristic PyTorch uses to choose between cuSOLVER and MAGMA for CUDA linear algebra operations. + + .. warning:: This flag is experimental and subject to change. + + When PyTorch runs a CUDA linear algebra operation it often uses the cuSOLVER or MAGMA libraries, + and if both are available it decides which to use with a heuristic. + This flag (a :class:`str`) allows overriding those heuristics. + + * If `"cusolver"` is set then cuSOLVER will be used wherever possible. + * If `"magma"` is set then MAGMA will be used wherever possible. + * If `"default"` (the default) is set then heuristics will be used to pick between + cuSOLVER and MAGMA if both are available. + * When no input is given, this function returns the currently preferred library. + * User may use the environment variable TORCH_LINALG_PREFER_CUSOLVER=1 to set the preferred library to cuSOLVER + globally. + This flag only sets the initial value of the preferred library and the preferred library + may still be overridden by this function call later in your script. + + Note: When a library is preferred other libraries may still be used if the preferred library + doesn't implement the operation(s) called. + This flag may achieve better performance if PyTorch's heuristic library selection is incorrect + for your application's inputs. + + Currently supported linalg operators: + + * :func:`torch.linalg.inv` + * :func:`torch.linalg.inv_ex` + * :func:`torch.linalg.cholesky` + * :func:`torch.linalg.cholesky_ex` + * :func:`torch.cholesky_solve` + * :func:`torch.cholesky_inverse` + * :func:`torch.linalg.lu_factor` + * :func:`torch.linalg.lu` + * :func:`torch.linalg.lu_solve` + * :func:`torch.linalg.qr` + * :func:`torch.linalg.eigh` + * :func:`torch.linalg.eighvals` + * :func:`torch.linalg.svd` + * :func:`torch.linalg.svdvals` + """ + if backend is None: + pass + elif isinstance(backend, str): + if backend not in _LinalgBackends: + raise RuntimeError( + "Unknown input value. " f"Choose from: {_LinalgBackends_str}." + ) + torch._C._set_linalg_preferred_backend(_LinalgBackends[backend]) + elif isinstance(backend, torch._C._LinalgBackend): + torch._C._set_linalg_preferred_backend(backend) + else: + raise RuntimeError("Unknown input value type.") + + return torch._C._get_linalg_preferred_backend() + + +_BlasBackends = { + "cublas": torch._C._BlasBackend.Cublas, + "cublaslt": torch._C._BlasBackend.Cublaslt, + "hipblaslt": torch._C._BlasBackend.Cublaslt, # alias + "ck": torch._C._BlasBackend.Ck, +} +_BlasBackends_str = ", ".join(_BlasBackends.keys()) + + +def preferred_blas_library( + backend: Union[None, str, torch._C._BlasBackend] = None +) -> torch._C._BlasBackend: + r""" + Override the library PyTorch uses for BLAS operations. Choose between cuBLAS, cuBLASLt, and CK [ROCm-only]. + + .. warning:: This flag is experimental and subject to change. + + When PyTorch runs a CUDA BLAS operation it defaults to cuBLAS even if both cuBLAS and cuBLASLt are available. + For PyTorch built for ROCm, hipBLAS, hipBLASLt, and CK may offer different performance. + This flag (a :class:`str`) allows overriding which BLAS library to use. + + * If `"cublas"` is set then cuBLAS will be used wherever possible. + * If `"cublaslt"` is set then cuBLASLt will be used wherever possible. + * If `"ck"` is set then CK will be used wherever possible. + * When no input is given, this function returns the currently preferred library. + * User may use the environment variable TORCH_BLAS_PREFER_CUBLASLT=1 to set the preferred library to cuBLASLt + globally. + This flag only sets the initial value of the preferred library and the preferred library + may still be overridden by this function call later in your script. + + Note: When a library is preferred other libraries may still be used if the preferred library + doesn't implement the operation(s) called. + This flag may achieve better performance if PyTorch's library selection is incorrect + for your application's inputs. + + """ + if backend is None: + pass + elif isinstance(backend, str): + if backend not in _BlasBackends: + raise RuntimeError( + "Unknown input value. " f"Choose from: {_BlasBackends_str}." + ) + torch._C._set_blas_preferred_backend(_BlasBackends[backend]) + elif isinstance(backend, torch._C._BlasBackend): + torch._C._set_blas_preferred_backend(backend) + else: + raise RuntimeError("Unknown input value type.") + + return torch._C._get_blas_preferred_backend() + + +from torch._C import _SDPAParams as SDPAParams, _SDPBackend as SDPBackend + + +# Set the __module__ attribute +SDPAParams.__module__ = "torch.backends.cuda" +SDPAParams.__name__ = "SDPAParams" + + +def flash_sdp_enabled(): + r""" + .. warning:: This flag is beta and subject to change. + + Returns whether flash scaled dot product attention is enabled or not. + """ + return torch._C._get_flash_sdp_enabled() + + +def enable_flash_sdp(enabled: bool): + r""" + .. warning:: This flag is beta and subject to change. + + Enables or disables flash scaled dot product attention. + """ + torch._C._set_sdp_use_flash(enabled) + + +def mem_efficient_sdp_enabled(): + r""" + .. warning:: This flag is beta and subject to change. + + Returns whether memory efficient scaled dot product attention is enabled or not. + """ + return torch._C._get_mem_efficient_sdp_enabled() + + +def enable_mem_efficient_sdp(enabled: bool): + r""" + .. warning:: This flag is beta and subject to change. + + Enables or disables memory efficient scaled dot product attention. + """ + torch._C._set_sdp_use_mem_efficient(enabled) + + +def math_sdp_enabled(): + r""" + .. warning:: This flag is beta and subject to change. + + Returns whether math scaled dot product attention is enabled or not. + """ + return torch._C._get_math_sdp_enabled() + + +def enable_math_sdp(enabled: bool): + r""" + .. warning:: This flag is beta and subject to change. + + Enables or disables math scaled dot product attention. + """ + torch._C._set_sdp_use_math(enabled) + + +def allow_fp16_bf16_reduction_math_sdp(enabled: bool): + r""" + .. warning:: This flag is beta and subject to change. + + Enables or disables fp16/bf16 reduction in math scaled dot product attention. + """ + torch._C._set_math_sdp_allow_fp16_bf16_reduction(enabled) + + +def fp16_bf16_reduction_math_sdp_allowed(): + r""" + .. warning:: This flag is beta and subject to change. + + Returns whether fp16/bf16 reduction in math scaled dot product attention is enabled or not. + """ + return torch._C._get_math_sdp_allow_fp16_bf16_reduction() + + +def is_flash_attention_available() -> bool: + r"""Check if PyTorch was built with FlashAttention for scaled_dot_product_attention. + + Returns: + True if FlashAttention is built and available; otherwise, False. + + Note: + This function is dependent on a CUDA-enabled build of PyTorch. It will return False + in non-CUDA environments. + """ + return torch._C._is_flash_attention_available() + + +def can_use_flash_attention(params: SDPAParams, debug: bool = False) -> bool: + r"""Check if FlashAttention can be utilized in scaled_dot_product_attention. + + Args: + params: An instance of SDPAParams containing the tensors for query, + key, value, an optional attention mask, dropout rate, and + a flag indicating if the attention is causal. + debug: Whether to logging.warn debug information as to why FlashAttention could not be run. + Defaults to False. + + Returns: + True if FlashAttention can be used with the given parameters; otherwise, False. + + Note: + This function is dependent on a CUDA-enabled build of PyTorch. It will return False + in non-CUDA environments. + """ + return torch._C._can_use_flash_attention(params, debug) + + +def can_use_efficient_attention(params: SDPAParams, debug: bool = False) -> bool: + r"""Check if efficient_attention can be utilized in scaled_dot_product_attention. + + Args: + params: An instance of SDPAParams containing the tensors for query, + key, value, an optional attention mask, dropout rate, and + a flag indicating if the attention is causal. + debug: Whether to logging.warn with information as to why efficient_attention could not be run. + Defaults to False. + + Returns: + True if efficient_attention can be used with the given parameters; otherwise, False. + + Note: + This function is dependent on a CUDA-enabled build of PyTorch. It will return False + in non-CUDA environments. + """ + return torch._C._can_use_mem_efficient_attention(params, debug) + + +def can_use_cudnn_attention(params: SDPAParams, debug: bool = False) -> bool: + r"""Check if cudnn_attention can be utilized in scaled_dot_product_attention. + + Args: + params: An instance of SDPAParams containing the tensors for query, + key, value, an optional attention mask, dropout rate, and + a flag indicating if the attention is causal. + debug: Whether to logging.warn with information as to why cuDNN attention could not be run. + Defaults to False. + + Returns: + True if cuDNN can be used with the given parameters; otherwise, False. + + Note: + This function is dependent on a CUDA-enabled build of PyTorch. It will return False + in non-CUDA environments. + """ + return torch._C._can_use_cudnn_attention(params, debug) + + +def cudnn_sdp_enabled(): + r""" + .. warning:: This flag is beta and subject to change. + + Returns whether cuDNN scaled dot product attention is enabled or not. + """ + return torch._C._get_cudnn_sdp_enabled() + + +def enable_cudnn_sdp(enabled: bool): + r""" + .. warning:: This flag is beta and subject to change. + + Enables or disables cuDNN scaled dot product attention. + """ + torch._C._set_sdp_use_cudnn(enabled) + + +@contextlib.contextmanager +@deprecated( + ( + "`torch.backends.cuda.sdp_kernel()` is deprecated. " + "In the future, this context manager will be removed. " + "Please see `torch.nn.attention.sdpa_kernel()` for the new context manager, " + "with updated signature." + ), + category=FutureWarning, +) +def sdp_kernel( + enable_flash: bool = True, + enable_math: bool = True, + enable_mem_efficient: bool = True, + enable_cudnn: bool = True, +): + r""" + .. warning:: This flag is beta and subject to change. + + This context manager can be used to temporarily enable or disable any of the three backends for scaled dot product attention. + Upon exiting the context manager, the previous state of the flags will be restored. + """ + from torch.nn.attention import sdpa_kernel + + backend_list = [] + if enable_flash: + backend_list.append(SDPBackend.FLASH_ATTENTION) + if enable_mem_efficient: + backend_list.append(SDPBackend.EFFICIENT_ATTENTION) + if enable_math: + backend_list.append(SDPBackend.MATH) + if enable_cudnn: + backend_list.append(SDPBackend.CUDNN_ATTENTION) + + with sdpa_kernel(backend_list) as context: + try: + yield context + finally: + pass + + +cufft_plan_cache = cuFFTPlanCacheManager() +matmul = cuBLASModule() diff --git a/lib/python3.12/site-packages/torch/backends/cuda/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/backends/cuda/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc4cb24fd5b21b838dff0518f7e7ee941ad99f45 Binary files /dev/null and b/lib/python3.12/site-packages/torch/backends/cuda/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/backends/cudnn/__init__.py b/lib/python3.12/site-packages/torch/backends/cudnn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0ba6f9e3b40bdfd8194dc9be97f542910ffb4632 --- /dev/null +++ b/lib/python3.12/site-packages/torch/backends/cudnn/__init__.py @@ -0,0 +1,208 @@ +# mypy: allow-untyped-defs +import os +import sys +import warnings +from contextlib import contextmanager +from typing import Optional + +import torch +from torch.backends import __allow_nonbracketed_mutation, ContextProp, PropModule + + +try: + from torch._C import _cudnn +except ImportError: + _cudnn = None # type: ignore[assignment] + +# Write: +# +# torch.backends.cudnn.enabled = False +# +# to globally disable CuDNN/MIOpen + +__cudnn_version: Optional[int] = None + +if _cudnn is not None: + + def _init(): + global __cudnn_version + if __cudnn_version is None: + __cudnn_version = _cudnn.getVersionInt() + runtime_version = _cudnn.getRuntimeVersion() + compile_version = _cudnn.getCompileVersion() + runtime_major, runtime_minor, _ = runtime_version + compile_major, compile_minor, _ = compile_version + # Different major versions are always incompatible + # Starting with cuDNN 7, minor versions are backwards-compatible + # Not sure about MIOpen (ROCm), so always do a strict check + if runtime_major != compile_major: + cudnn_compatible = False + elif runtime_major < 7 or not _cudnn.is_cuda: + cudnn_compatible = runtime_minor == compile_minor + else: + cudnn_compatible = runtime_minor >= compile_minor + if not cudnn_compatible: + if os.environ.get("PYTORCH_SKIP_CUDNN_COMPATIBILITY_CHECK", "0") == "1": + return True + base_error_msg = ( + f"cuDNN version incompatibility: " + f"PyTorch was compiled against {compile_version} " + f"but found runtime version {runtime_version}. " + f"PyTorch already comes bundled with cuDNN. " + f"One option to resolving this error is to ensure PyTorch " + f"can find the bundled cuDNN. " + ) + + if "LD_LIBRARY_PATH" in os.environ: + ld_library_path = os.environ.get("LD_LIBRARY_PATH", "") + if any( + substring in ld_library_path for substring in ["cuda", "cudnn"] + ): + raise RuntimeError( + f"{base_error_msg}" + f"Looks like your LD_LIBRARY_PATH contains incompatible version of cudnn. " + f"Please either remove it from the path or install cudnn {compile_version}" + ) + else: + raise RuntimeError( + f"{base_error_msg}" + f"one possibility is that there is a " + f"conflicting cuDNN in LD_LIBRARY_PATH." + ) + else: + raise RuntimeError(base_error_msg) + + return True + +else: + + def _init(): + return False + + +def version(): + """Return the version of cuDNN.""" + if not _init(): + return None + return __cudnn_version + + +CUDNN_TENSOR_DTYPES = { + torch.half, + torch.float, + torch.double, +} + + +def is_available(): + r"""Return a bool indicating if CUDNN is currently available.""" + return torch._C._has_cudnn + + +def is_acceptable(tensor): + if not torch._C._get_cudnn_enabled(): + return False + if tensor.device.type != "cuda" or tensor.dtype not in CUDNN_TENSOR_DTYPES: + return False + if not is_available(): + warnings.warn( + "PyTorch was compiled without cuDNN/MIOpen support. To use cuDNN/MIOpen, rebuild " + "PyTorch making sure the library is visible to the build system." + ) + return False + if not _init(): + warnings.warn( + "cuDNN/MIOpen library not found. Check your {libpath}".format( + libpath={"darwin": "DYLD_LIBRARY_PATH", "win32": "PATH"}.get( + sys.platform, "LD_LIBRARY_PATH" + ) + ) + ) + return False + return True + + +def set_flags( + _enabled=None, + _benchmark=None, + _benchmark_limit=None, + _deterministic=None, + _allow_tf32=None, +): + orig_flags = ( + torch._C._get_cudnn_enabled(), + torch._C._get_cudnn_benchmark(), + None if not is_available() else torch._C._cuda_get_cudnn_benchmark_limit(), + torch._C._get_cudnn_deterministic(), + torch._C._get_cudnn_allow_tf32(), + ) + if _enabled is not None: + torch._C._set_cudnn_enabled(_enabled) + if _benchmark is not None: + torch._C._set_cudnn_benchmark(_benchmark) + if _benchmark_limit is not None and is_available(): + torch._C._cuda_set_cudnn_benchmark_limit(_benchmark_limit) + if _deterministic is not None: + torch._C._set_cudnn_deterministic(_deterministic) + if _allow_tf32 is not None: + torch._C._set_cudnn_allow_tf32(_allow_tf32) + return orig_flags + + +@contextmanager +def flags( + enabled=False, + benchmark=False, + benchmark_limit=10, + deterministic=False, + allow_tf32=True, +): + with __allow_nonbracketed_mutation(): + orig_flags = set_flags( + enabled, benchmark, benchmark_limit, deterministic, allow_tf32 + ) + try: + yield + finally: + # recover the previous values + with __allow_nonbracketed_mutation(): + set_flags(*orig_flags) + + +# The magic here is to allow us to intercept code like this: +# +# torch.backends..enabled = True + + +class CudnnModule(PropModule): + def __init__(self, m, name): + super().__init__(m, name) + + enabled = ContextProp(torch._C._get_cudnn_enabled, torch._C._set_cudnn_enabled) + deterministic = ContextProp( + torch._C._get_cudnn_deterministic, torch._C._set_cudnn_deterministic + ) + benchmark = ContextProp( + torch._C._get_cudnn_benchmark, torch._C._set_cudnn_benchmark + ) + benchmark_limit = None + if is_available(): + benchmark_limit = ContextProp( + torch._C._cuda_get_cudnn_benchmark_limit, + torch._C._cuda_set_cudnn_benchmark_limit, + ) + allow_tf32 = ContextProp( + torch._C._get_cudnn_allow_tf32, torch._C._set_cudnn_allow_tf32 + ) + + +# This is the sys.modules replacement trick, see +# https://stackoverflow.com/questions/2447353/getattr-on-a-module/7668273#7668273 +sys.modules[__name__] = CudnnModule(sys.modules[__name__], __name__) + +# Add type annotation for the replaced module +enabled: bool +deterministic: bool +benchmark: bool +allow_tf32: bool +benchmark_limit: int diff --git a/lib/python3.12/site-packages/torch/backends/cudnn/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/backends/cudnn/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..114b8c206e046b372f1bdab83d37b2ccee11d169 Binary files /dev/null and b/lib/python3.12/site-packages/torch/backends/cudnn/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/backends/cudnn/__pycache__/rnn.cpython-312.pyc b/lib/python3.12/site-packages/torch/backends/cudnn/__pycache__/rnn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8086b33297ed47d252795d8f23cfec5dfd3988c Binary files /dev/null and b/lib/python3.12/site-packages/torch/backends/cudnn/__pycache__/rnn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/backends/cudnn/rnn.py b/lib/python3.12/site-packages/torch/backends/cudnn/rnn.py new file mode 100644 index 0000000000000000000000000000000000000000..5b253e19054027e88c3ead9e536d0fa0cd0da3cc --- /dev/null +++ b/lib/python3.12/site-packages/torch/backends/cudnn/rnn.py @@ -0,0 +1,64 @@ +# mypy: allow-untyped-defs +import torch.cuda + + +try: + from torch._C import _cudnn +except ImportError: + # Uses of all the functions below should be guarded by torch.backends.cudnn.is_available(), + # so it's safe to not emit any checks here. + _cudnn = None # type: ignore[assignment] + + +def get_cudnn_mode(mode): + if mode == "RNN_RELU": + return int(_cudnn.RNNMode.rnn_relu) + elif mode == "RNN_TANH": + return int(_cudnn.RNNMode.rnn_tanh) + elif mode == "LSTM": + return int(_cudnn.RNNMode.lstm) + elif mode == "GRU": + return int(_cudnn.RNNMode.gru) + else: + raise Exception(f"Unknown mode: {mode}") # noqa: TRY002 + + +# NB: We don't actually need this class anymore (in fact, we could serialize the +# dropout state for even better reproducibility), but it is kept for backwards +# compatibility for old models. +class Unserializable: + def __init__(self, inner): + self.inner = inner + + def get(self): + return self.inner + + def __getstate__(self): + # Note: can't return {}, because python2 won't call __setstate__ + # if the value evaluates to False + return "" + + def __setstate__(self, state): + self.inner = None + + +def init_dropout_state(dropout, train, dropout_seed, dropout_state): + dropout_desc_name = "desc_" + str(torch.cuda.current_device()) + dropout_p = dropout if train else 0 + if (dropout_desc_name not in dropout_state) or ( + dropout_state[dropout_desc_name].get() is None + ): + if dropout_p == 0: + dropout_state[dropout_desc_name] = Unserializable(None) + else: + dropout_state[dropout_desc_name] = Unserializable( + torch._cudnn_init_dropout_state( # type: ignore[call-arg] + dropout_p, + train, + dropout_seed, + self_ty=torch.uint8, + device=torch.device("cuda"), + ) + ) + dropout_ts = dropout_state[dropout_desc_name].get() + return dropout_ts diff --git a/lib/python3.12/site-packages/torch/backends/cusparselt/__init__.py b/lib/python3.12/site-packages/torch/backends/cusparselt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..da46274a2846da188d02f9b5b66a7dbd7d2cd38f --- /dev/null +++ b/lib/python3.12/site-packages/torch/backends/cusparselt/__init__.py @@ -0,0 +1,57 @@ +# mypy: allow-untyped-defs +from typing import Optional + +import torch + + +__all__ = [ + "version", + "is_available", + "get_max_alg_id", +] + +try: + from torch._C import _cusparselt +except ImportError: + _cusparselt = None # type: ignore[assignment] + +__cusparselt_version: Optional[int] = None +__MAX_ALG_ID: Optional[int] = None + +if _cusparselt is not None: + + def _init(): + global __cusparselt_version + global __MAX_ALG_ID + if __cusparselt_version is None: + __cusparselt_version = _cusparselt.getVersionInt() + if __cusparselt_version == 400: + __MAX_ALG_ID = 4 + elif __cusparselt_version == 502: + __MAX_ALG_ID = 5 + elif __cusparselt_version == 602: + __MAX_ALG_ID = 37 + return True + +else: + + def _init(): + return False + + +def version() -> Optional[int]: + """Return the version of cuSPARSELt""" + if not _init(): + return None + return __cusparselt_version + + +def is_available() -> bool: + r"""Return a bool indicating if cuSPARSELt is currently available.""" + return torch._C._has_cusparselt + + +def get_max_alg_id() -> Optional[int]: + if not _init(): + return None + return __MAX_ALG_ID diff --git a/lib/python3.12/site-packages/torch/backends/cusparselt/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/backends/cusparselt/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e92d3bc8470f9ef829309902bedef0c0131de950 Binary files /dev/null and b/lib/python3.12/site-packages/torch/backends/cusparselt/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/backends/mkldnn/__init__.py b/lib/python3.12/site-packages/torch/backends/mkldnn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8d4b9c49519b107dac133c331073bb61865896f6 --- /dev/null +++ b/lib/python3.12/site-packages/torch/backends/mkldnn/__init__.py @@ -0,0 +1,100 @@ +# mypy: allow-untyped-defs +import sys +from contextlib import contextmanager +from typing import TYPE_CHECKING + +import torch +from torch.backends import __allow_nonbracketed_mutation, ContextProp, PropModule + + +def is_available(): + r"""Return whether PyTorch is built with MKL-DNN support.""" + return torch._C._has_mkldnn + + +VERBOSE_OFF = 0 +VERBOSE_ON = 1 +VERBOSE_ON_CREATION = 2 + + +class verbose: + """ + On-demand oneDNN (former MKL-DNN) verbosing functionality. + + To make it easier to debug performance issues, oneDNN can dump verbose + messages containing information like kernel size, input data size and + execution duration while executing the kernel. The verbosing functionality + can be invoked via an environment variable named `DNNL_VERBOSE`. However, + this methodology dumps messages in all steps. Those are a large amount of + verbose messages. Moreover, for investigating the performance issues, + generally taking verbose messages for one single iteration is enough. + This on-demand verbosing functionality makes it possible to control scope + for verbose message dumping. In the following example, verbose messages + will be dumped out for the second inference only. + + .. highlight:: python + .. code-block:: python + + import torch + model(data) + with torch.backends.mkldnn.verbose(torch.backends.mkldnn.VERBOSE_ON): + model(data) + + Args: + level: Verbose level + - ``VERBOSE_OFF``: Disable verbosing + - ``VERBOSE_ON``: Enable verbosing + - ``VERBOSE_ON_CREATION``: Enable verbosing, including oneDNN kernel creation + """ + + def __init__(self, level): + self.level = level + + def __enter__(self): + if self.level == VERBOSE_OFF: + return + st = torch._C._verbose.mkldnn_set_verbose(self.level) + assert ( + st + ), "Failed to set MKLDNN into verbose mode. Please consider to disable this verbose scope." + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + torch._C._verbose.mkldnn_set_verbose(VERBOSE_OFF) + return False + + +def set_flags(_enabled, _deterministic=None): + orig_flags = (torch._C._get_mkldnn_enabled(), torch._C._get_mkldnn_deterministic()) + torch._C._set_mkldnn_enabled(_enabled) + if _deterministic is not None: + torch._C._set_mkldnn_deterministic(_deterministic) + return orig_flags + + +@contextmanager +def flags(enabled=False, deterministic=False): + with __allow_nonbracketed_mutation(): + orig_flags = set_flags(enabled, deterministic) + try: + yield + finally: + with __allow_nonbracketed_mutation(): + set_flags(*orig_flags) + + +class MkldnnModule(PropModule): + def __init__(self, m, name): + super().__init__(m, name) + + enabled = ContextProp(torch._C._get_mkldnn_enabled, torch._C._set_mkldnn_enabled) + deterministic = ContextProp( + torch._C._get_mkldnn_deterministic, torch._C._set_mkldnn_deterministic + ) + + +if TYPE_CHECKING: + enabled: ContextProp + deterministic: ContextProp + +sys.modules[__name__] = MkldnnModule(sys.modules[__name__], __name__) diff --git a/lib/python3.12/site-packages/torch/backends/mkldnn/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/backends/mkldnn/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8945d10b804d8f9e5538fae4b2a8687c6a7de5a1 Binary files /dev/null and b/lib/python3.12/site-packages/torch/backends/mkldnn/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/backends/openmp/__init__.py b/lib/python3.12/site-packages/torch/backends/openmp/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aff8d46cd4ac2d9ff49942542d99ac2afbb85896 --- /dev/null +++ b/lib/python3.12/site-packages/torch/backends/openmp/__init__.py @@ -0,0 +1,7 @@ +# mypy: allow-untyped-defs +import torch + + +def is_available(): + r"""Return whether PyTorch is built with OpenMP support.""" + return torch._C.has_openmp diff --git a/lib/python3.12/site-packages/torch/backends/openmp/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/backends/openmp/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c10021cec239823d8b378003b4d4124bb06f598 Binary files /dev/null and b/lib/python3.12/site-packages/torch/backends/openmp/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..27ad354d5ea5df1d85e261c29b11c1fd301d7911 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_async.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_async.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..562bf50098632d8a26039b5c149429fcc4a9688c Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_async.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_await.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_await.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..299e4ca89307c8718183ab1f446382872c077522 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_await.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_builtins.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_builtins.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2bed41f92aac75e2b52514e0dd72711f47ee9759 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_builtins.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_check.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_check.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b809b272fb4efc057098aee8af220f35b6357933 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_check.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_dataclass_impls.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_dataclass_impls.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..981de64d76f552304fcaaf7e063b84aaa3682090 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_dataclass_impls.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_decomposition_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_decomposition_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0845ba2cf4b51bd000bc875620964e6532d94c9 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_decomposition_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_decompositions.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_decompositions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8b534f0eacd2ed18b1a1b072a0669cdbcc088c3 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_decompositions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_freeze.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_freeze.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3cd7fc46c817dbe237e7303548a5889ca52477f4 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_freeze.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_fuser.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_fuser.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b59ddc32a273850036a17153fc4a8564c587e0f Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_fuser.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_ir_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_ir_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e5e15625a5c8868bb39f87673c82e88b7e6bd20 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_ir_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_logging.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c7ffac0c12899ad9ae034522fa0e4e76d7801df Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_logging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_monkeytype_config.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_monkeytype_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a93c961530989db4ad53934a2744dd0238f893fd Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_monkeytype_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_pickle.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_pickle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5cfcd821993ea31951bda7d4a01d34fba9d3e2f1 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_pickle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_recursive.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_recursive.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d53e0086e7698de9062a3fc93c22883eab426b36 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_recursive.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_script.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_script.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a074a25d1cbdf5ff9ab100a7e24fee8a5eb56323 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_script.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_serialization.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_serialization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bec0529039894928ee92c641ebbfb84fb6f53c3b Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_serialization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_shape_functions.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_shape_functions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7addac195ff445358bc3b2f66c51a2285cf6082 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_shape_functions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_state.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_state.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..303018b03d5e8464e5c02a00522cc47f91d6ab69 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_state.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/_trace.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/_trace.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47793d505d5b9e587142943af83780f082e73a4b Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/_trace.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/annotations.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/annotations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdbe588992bcdf67c4cd189001eceb2fe7fdc0c3 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/annotations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/frontend.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/frontend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a379f2ccd3274619ff571ac73180b6ca6e6aee67 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/frontend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/generate_bytecode.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/generate_bytecode.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e20eec749b719705ad2037493d0d015cd2d583d2 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/generate_bytecode.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/quantized.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/quantized.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee14ccf738cbfdb4c538e598f9efda3230242711 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/quantized.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/supported_ops.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/supported_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8d805322bb2a506422b463a6419b3fcacdabc07 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/supported_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/__pycache__/unsupported_tensor_ops.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/__pycache__/unsupported_tensor_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68ae74b13c0427b8d2d8269e1f3d1eda30fb2c69 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/__pycache__/unsupported_tensor_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/_async.py b/lib/python3.12/site-packages/torch/jit/_async.py new file mode 100644 index 0000000000000000000000000000000000000000..9ccde01d4a425632fa942df18300110620c0248a --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_async.py @@ -0,0 +1,102 @@ +# mypy: allow-untyped-defs +"""Async API. + +This module contains the API for parallelism in TorchScript, notably: + * torch.jit.fork + * torch.jit.wait + +This is not intended to be imported directly; please use the exposed +functionalities in `torch.jit`. +""" + +import torch +from torch._jit_internal import Future +from torch.jit._builtins import _register_builtin +from torch.utils import set_module + + +set_module(Future, "torch.jit") + + +def fork(func, *args, **kwargs): + r""" + Create an asynchronous task executing `func` and a reference to the value of the result of this execution. + + `fork` will return immediately, so the return value of `func` may not have been computed yet. To force completion + of the task and access the return value invoke `torch.jit.wait` on the Future. `fork` invoked + with a `func` which returns `T` is typed as `torch.jit.Future[T]`. `fork` calls can be arbitrarily + nested, and may be invoked with positional and keyword arguments. + Asynchronous execution will only occur when run in TorchScript. If run in pure python, + `fork` will not execute in parallel. `fork` will also not execute in parallel when invoked + while tracing, however the `fork` and `wait` calls will be captured in the exported IR Graph. + + .. warning:: + `fork` tasks will execute non-deterministically. We recommend only spawning + parallel fork tasks for pure functions that do not modify their inputs, + module attributes, or global state. + + Args: + func (callable or torch.nn.Module): A Python function or `torch.nn.Module` + that will be invoked. If executed in TorchScript, it will execute asynchronously, + otherwise it will not. Traced invocations of fork will be captured in the IR. + ``*args``, ``**kwargs``: arguments to invoke `func` with. + Returns: + `torch.jit.Future[T]`: a reference to the execution of `func`. The value `T` + can only be accessed by forcing completion of `func` through `torch.jit.wait`. + + Example (fork a free function): + + .. code-block:: python + + import torch + from torch import Tensor + def foo(a : Tensor, b : int) -> Tensor: + return a + b + def bar(a): + fut : torch.jit.Future[Tensor] = torch.jit.fork(foo, a, b=2) + return torch.jit.wait(fut) + script_bar = torch.jit.script(bar) + input = torch.tensor(2) + # only the scripted version executes asynchronously + assert script_bar(input) == bar(input) + # trace is not run asynchronously, but fork is captured in IR + graph = torch.jit.trace(bar, (input,)).graph + assert "fork" in str(graph) + + Example (fork a module method): + + .. code-block:: python + + import torch + from torch import Tensor + class AddMod(torch.nn.Module): + def forward(self, a: Tensor, b : int): + return a + b + class Mod(torch.nn.Module): + def __init__(self) -> None: + super(self).__init__() + self.mod = AddMod() + def forward(self, input): + fut = torch.jit.fork(self.mod, a, b=2) + return torch.jit.wait(fut) + input = torch.tensor(2) + mod = Mod() + assert mod(input) == torch.jit.script(mod).forward(input) + """ + return torch._C.fork(func, *args, **kwargs) + + +def wait(future): + r""" + Force completion of a `torch.jit.Future[T]` asynchronous task, returning the result of the task. + + See :func:`~fork` for docs and examples. + Args: + future (torch.jit.Future[T]): an asynchronous task reference, created through `torch.jit.fork` + Returns: + `T`: the return value of the completed task + """ + return torch._C.wait(future) + + +_register_builtin(wait, "aten::wait") diff --git a/lib/python3.12/site-packages/torch/jit/_await.py b/lib/python3.12/site-packages/torch/jit/_await.py new file mode 100644 index 0000000000000000000000000000000000000000..b00da7b3384b7a263aeab613b3e613c4bd3c3667 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_await.py @@ -0,0 +1,27 @@ +# mypy: allow-untyped-defs +import torch +from torch._jit_internal import _Await +from torch.jit._builtins import _register_builtin +from torch.utils import set_module + + +set_module(_Await, "torch.jit") + + +def _awaitable(func, *args, **kwargs): + r"""Create Await object that will call specified functioni with specified args, when it is requested for the result.""" + return torch._C._awaitable(func, *args, **kwargs) + + +def _awaitable_wait(aw): + r"""Request await the result of execution, if Await is not completed yet, the func will be called immediately.""" + return torch._C._awaitable_wait(aw) + + +def _awaitable_nowait(o): + r"""Create completed Await with specified result.""" + return torch._C._awaitable_nowait(o) + + +_register_builtin(_awaitable_wait, "prim::awaitable_wait") +_register_builtin(_awaitable_nowait, "prim::awaitable_nowait") diff --git a/lib/python3.12/site-packages/torch/jit/_builtins.py b/lib/python3.12/site-packages/torch/jit/_builtins.py new file mode 100644 index 0000000000000000000000000000000000000000..a3edc688d9e14b33eccc146c91f171f90af8ab26 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_builtins.py @@ -0,0 +1,193 @@ +# mypy: allow-untyped-defs +import cmath +import math +import warnings +from collections import OrderedDict +from typing import Dict, Optional + +import torch +import torch.backends.cudnn as cudnn +from torch.nn.modules.utils import ( + _list_with_default, + _pair, + _quadruple, + _single, + _triple, +) + + +_builtin_table: Optional[Dict[int, str]] = None + +_modules_containing_builtins = (torch, torch._C._nn, torch._C._fft, torch._C._linalg, torch._C._nested, torch._C._sparse, torch._C._special) # type: ignore[attr-defined] # noqa: B950 + +_builtin_ops = [ + # Pairs of (function, op_name) + (_pair, "aten::_pair"), + (_quadruple, "aten::_quadruple"), + (_single, "aten::_single"), + (_triple, "aten::_triple"), + (_list_with_default, "aten::list_with_default"), + (OrderedDict, "aten::dict"), + (dict, "aten::dict"), + (cudnn.is_acceptable, "aten::cudnn_is_acceptable"), + (math.ceil, "aten::ceil"), + (math.copysign, "aten::copysign"), + (math.erf, "aten::erf"), + (math.erfc, "aten::erfc"), + (math.exp, "aten::exp"), + (math.expm1, "aten::expm1"), + (math.fabs, "aten::fabs"), + (math.floor, "aten::floor"), + (math.gamma, "aten::gamma"), + (math.lgamma, "aten::lgamma"), + (math.log, "aten::log"), + (math.log10, "aten::log10"), + (math.log1p, "aten::log1p"), + (math.pow, "aten::pow"), + (math.sqrt, "aten::sqrt"), + (math.isnan, "aten::isnan"), + (math.asinh, "aten::asinh"), + (math.atanh, "aten::atanh"), + (math.cosh, "aten::cosh"), + (math.sinh, "aten::sinh"), + (math.tanh, "aten::tanh"), + (math.acos, "aten::acos"), + (math.asin, "aten::asin"), + (math.atan, "aten::atan"), + (math.atan2, "aten::atan2"), + (math.cos, "aten::cos"), + (math.sin, "aten::sin"), + (math.tan, "aten::tan"), + (math.asinh, "aten::asinh"), + (math.atanh, "aten::atanh"), + (math.acosh, "aten::acosh"), + (math.fmod, "aten::fmod"), + (math.modf, "aten::modf"), + (math.factorial, "aten::factorial"), + (math.frexp, "aten::frexp"), + (math.isinf, "aten::isinf"), + (math.degrees, "aten::degrees"), + (math.radians, "aten::radians"), + (cmath.isnan, "aten::isnan"), + (cmath.isfinite, "aten::isfinite"), + (cmath.isinf, "aten::isinf"), + (cmath.phase, "aten::angle"), + (cmath.rect, "aten::polar"), + (cmath.log, "aten::log"), + (cmath.log10, "aten::log10"), + (cmath.sqrt, "aten::sqrt"), + (cmath.exp, "aten::exp"), + (cmath.sin, "aten::sin"), + (cmath.tan, "aten::tan"), + (cmath.cos, "aten::cos"), + (cmath.asin, "aten::asin"), + (cmath.acos, "aten::acos"), + (cmath.atan, "aten::atan"), + (cmath.sinh, "aten::sinh"), + (cmath.cosh, "aten::cosh"), + (cmath.tanh, "aten::tanh"), + (cmath.asinh, "aten::asinh"), + (cmath.acosh, "aten::acosh"), + (cmath.atanh, "aten::atanh"), + (math.ldexp, "aten::ldexp"), + (torch._assert, "aten::_assert"), + (torch.autograd.grad, "aten::grad"), + (torch.autograd.backward, "aten::backward"), + (torch._C._infer_size, "aten::_infer_size"), + (torch.nn.functional._no_grad_embedding_renorm_, "aten::_no_grad_embedding_renorm_"), # type: ignore[attr-defined] + (torch.nn.functional.assert_int_or_pair, "aten::_assert_int_or_pair"), + (torch.nn.init._no_grad_fill_, "aten::_no_grad_fill_"), + (torch.nn.init._no_grad_normal_, "aten::_no_grad_normal_"), + (torch.nn.init._no_grad_uniform_, "aten::_no_grad_uniform_"), + (torch.nn.init._no_grad_zero_, "aten::_no_grad_zero_"), + (torch._C._get_tracing_state, "aten::_get_tracing_state"), + (torch._C._get_cpu_capability, "aten::_get_cpu_capability"), + (warnings.warn, "aten::warn"), + (torch._VF.stft, "aten::stft"), # type: ignore[attr-defined] + (torch._VF.istft, "aten::istft"), # type: ignore[attr-defined] + (torch._VF.cdist, "aten::cdist"), # type: ignore[attr-defined] + (torch._VF.norm, "aten::norm"), # type: ignore[attr-defined] + (torch._VF.unique_dim, "aten::unique_dim"), + (torch._VF.unique_consecutive, "aten::unique_consecutive"), # type: ignore[attr-defined] + (torch._VF.nuclear_norm, "aten::nuclear_norm"), + (torch._VF.frobenius_norm, "aten::frobenius_norm"), + (torch._VF.tensordot, "aten::tensordot"), # type: ignore[attr-defined] +] + +# ops in torch.functional are bound to torch +# in these cases, we want to resolve the function to their python implementation +# instead looking up a builtin "aten::" schema + + +def _gen_torch_functional_registered_ops(): + # eventually ops should encompass all of torch/functional.py, (torch.functional.__all__) + # but we are currently only able to compile some of the functions. additionally, + # some functions directly map to their aten:: implementations. + # TODO: add support for more ops + ops = [ + "stft", + "istft", + "lu", + "cdist", + "norm", + "unique", + "unique_consecutive", + "tensordot", + ] + return {getattr(torch.functional, name) for name in ops} + + +_functional_registered_ops = _gen_torch_functional_registered_ops() + + +def _is_special_functional_bound_op(fn): + return fn in _functional_registered_ops + + +# lazily built to ensure the correct initialization order +def _get_builtin_table(): + global _builtin_table + if _builtin_table is not None: + return _builtin_table + _builtin_table = {} + + def register_all(mod): + for name in dir(mod): + v = getattr(mod, name) + if ( + callable(v) + and not _is_special_functional_bound_op(v) + and v is not torch.no_grad + and v is not torch.autocast + ): + # Fixup inconsistency in segment_reduce + if name == "_segment_reduce": + name = name[1:] + _builtin_ops.append((v, "aten::" + name)) + + for mod in _modules_containing_builtins: + register_all(mod) + + _builtin_ops.append((math.gcd, "aten::gcd")) + _builtin_ops.append((math.isfinite, "aten::isfinite")) + _builtin_ops.append((math.remainder, "aten::mathremainder")) # type: ignore[attr-defined] + + import torch.distributed.autograd as dist_autograd + + if dist_autograd.is_available(): + _builtin_ops.append((dist_autograd.get_gradients, "aten::get_gradients")) + _builtin_ops.append((dist_autograd.backward, "aten::dist_backward")) + + # populate the _builtin_table from _builtin_ops + for builtin, aten_op in _builtin_ops: + _builtin_table[id(builtin)] = aten_op + + return _builtin_table + + +def _register_builtin(fn, op): + _get_builtin_table()[id(fn)] = op + + +def _find_builtin(fn): + return _get_builtin_table().get(id(fn)) diff --git a/lib/python3.12/site-packages/torch/jit/_decomposition_utils.py b/lib/python3.12/site-packages/torch/jit/_decomposition_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..795f9da8e073a1197f81ddf379738a7f93605d7c --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_decomposition_utils.py @@ -0,0 +1,12 @@ +# mypy: allow-untyped-defs +import torch +from torch._ops import OpOverload, OpOverloadPacket + + +def _register_decomposition(op: OpOverload, graph: torch._C.Graph): + assert not isinstance( + op, OpOverloadPacket + ), f"Must pass specific op overload, not overload packet, found {op}" + assert isinstance(op, OpOverload) + + torch._C._jit_register_decomposition_for_schema(op._schema, graph) diff --git a/lib/python3.12/site-packages/torch/jit/_freeze.py b/lib/python3.12/site-packages/torch/jit/_freeze.py new file mode 100644 index 0000000000000000000000000000000000000000..e496bd74762554e4909362e8eebcb2711be04789 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_freeze.py @@ -0,0 +1,228 @@ +# mypy: allow-untyped-defs +"""Freezing. + +This is not intended to be imported directly; please use the exposed +functionalities in `torch.jit`. +""" + +from typing import List, Optional + +import torch +from torch.jit._script import RecursiveScriptModule, ScriptModule + + +def freeze( + mod, preserved_attrs: Optional[List[str]] = None, optimize_numerics: bool = True +): + r"""Freeze ScriptModule, inline submodules, and attributes as constants. + + Freezing a :class:`ScriptModule` will clone it and attempt to inline the cloned + module's submodules, parameters, and attributes as constants in the TorchScript IR Graph. + By default, `forward` will be preserved, as well as attributes & methods specified in + `preserved_attrs`. Additionally, any attribute that is modified within a preserved + method will be preserved. + + Freezing currently only accepts ScriptModules that are in eval mode. + + Freezing applies generic optimization that will speed up your model regardless of machine. + To further optimize using server-specific settings, run `optimize_for_inference` after + freezing. + + Args: + mod (:class:`ScriptModule`): a module to be frozen + preserved_attrs (Optional[List[str]]): a list of attributes to preserve in addition to the forward method. + Attributes modified in preserved methods will also be preserved. + optimize_numerics (bool): If ``True``, a set of optimization passes will be run that does not strictly + preserve numerics. Full details of optimization can be found at `torch.jit.run_frozen_optimizations`. + + Returns: + Frozen :class:`ScriptModule`. + + Example (Freezing a simple module with a Parameter): + + .. testcode:: + import torch + class MyModule(torch.nn.Module): + def __init__(self, N, M): + super().__init__() + self.weight = torch.nn.Parameter(torch.rand(N, M)) + self.linear = torch.nn.Linear(N, M) + + def forward(self, input): + output = self.weight.mm(input) + output = self.linear(output) + return output + + scripted_module = torch.jit.script(MyModule(2, 3).eval()) + frozen_module = torch.jit.freeze(scripted_module) + # parameters have been removed and inlined into the Graph as constants + assert len(list(frozen_module.named_parameters())) == 0 + # See the compiled graph as Python code + print(frozen_module.code) + + Example (Freezing a module with preserved attributes) + + .. testcode:: + import torch + class MyModule2(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.modified_tensor = torch.tensor(10.) + self.version = 1 + + def forward(self, input): + self.modified_tensor += 1 + return input + self.modified_tensor + + scripted_module = torch.jit.script(MyModule2().eval()) + frozen_module = torch.jit.freeze(scripted_module, preserved_attrs=["version"]) + # we've manually preserved `version`, so it still exists on the frozen module and can be modified + assert frozen_module.version == 1 + frozen_module.version = 2 + # `modified_tensor` is detected as being mutated in the forward, so freezing preserves + # it to retain model semantics + assert frozen_module(torch.tensor(1)) == torch.tensor(12) + # now that we've run it once, the next result will be incremented by one + assert frozen_module(torch.tensor(1)) == torch.tensor(13) + + Note: + Freezing submodule attributes is also supported: + frozen_module = torch.jit.freeze(scripted_module, preserved_attrs=["submodule.version"]) + + Note: + If you're not sure why an attribute is not being inlined as a constant, you can run + `dump_alias_db` on frozen_module.forward.graph to see if freezing has detected the + attribute is being modified. + + Note: + Because freezing makes weights constants and removes module hierarchy, `to` and other + nn.Module methods to manipulate device or dtype no longer work. As a workaround, + You can remap devices by specifying `map_location` in `torch.jit.load`, however + device-specific logic may have been baked into the model. + """ + if not isinstance(mod, ScriptModule): + raise RuntimeError( + "Freezing expects a ScriptModule as input. " + "Please use torch.jit.script or torch.jit.trace to script your 'nn.Module'." + ) + + if mod.training: + raise RuntimeError( + "Freezing is currently only implemented for modules in eval mode. " + "Please call .eval() on your module before freezing." + ) + + preserved_attrs = preserved_attrs if preserved_attrs is not None else [] + + out = RecursiveScriptModule(torch._C._freeze_module(mod._c, preserved_attrs)) + RecursiveScriptModule._finalize_scriptmodule(out) + + preserved_methods = [x for x in preserved_attrs if mod._c._has_method(x)] + run_frozen_optimizations(out, optimize_numerics, preserved_methods) + + return out + + +def run_frozen_optimizations( + mod, optimize_numerics: bool = True, preserved_methods: Optional[List[str]] = None +): + r""" + Run a series of optimizations looking for patterns that occur in frozen graphs. + + The current set of optimizations includes: + - Dropout Removal + - Pretranspose Linear Layers + - Concat Linear Layers with same input Tensor + - Conv -> Batchnorm folding + - Conv -> Add/Sub folding + - Conv -> Mul/Div folding + + Args: + mod (:class:`ScriptModule`): a frozen module to be optimized + + optimize_numerics (bool): If ``True``, a set of optimization passes will be run that does not strictly + preserve numerics. These optimizations preserve default rtol and atol of `torch.testing.assert_close` + when applied on a single transformation, however in a module where many transformations are applied + the rtol or atol may no longer fall within the default `assert_close` tolerance. Conv -> Batchnorm folding, + Conv-Add/Sub, and Conv -> Mul/Div folding all may alter numerics. + + Returns: + None + + Note: + In rare occassions, this can result in slower execution. + + Example (Freezing a module with Conv->Batchnorm) + .. code-block:: python + import torch + in_channels, out_channels = 3, 32 + conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=True) + bn = torch.nn.BatchNorm2d(out_channels, eps=.001) + mod = torch.nn.Sequential(conv, bn) + # set optimize to False here, by default freezing runs run_frozen_optimizations + frozen_mod = torch.jit.freeze(torch.jit.script(mod.eval()), optimize=False) + # inspect frozen mod + assert "batch_norm" in str(frozen_mod.graph) + torch.jit.run_frozen_optimizations(frozen_mod) + assert "batch_norm" not in str(frozen_mod.graph) + + """ + if mod._c._has_method("forward"): + torch._C._jit_pass_optimize_frozen_graph(mod.graph, optimize_numerics) + + if preserved_methods is None: + preserved_methods = [] + + for method in preserved_methods: + torch._C._jit_pass_optimize_frozen_graph( + mod.__getattr__(method).graph, optimize_numerics + ) + + +def optimize_for_inference( + mod: ScriptModule, other_methods: Optional[List[str]] = None +) -> ScriptModule: + """ + Perform a set of optimization passes to optimize a model for the purposes of inference. + + If the model is not already frozen, optimize_for_inference + will invoke `torch.jit.freeze` automatically. + + In addition to generic optimizations that should speed up your model regardless + of environment, prepare for inference will also bake in build specific settings + such as the presence of CUDNN or MKLDNN, and may in the future make transformations + which speed things up on one machine but slow things down on another. Accordingly, + serialization is not implemented following invoking `optimize_for_inference` and + is not guaranteed. + + This is still in prototype, and may have the potential to slow down your model. + Primary use cases that have been targeted so far have been vision models on cpu + and gpu to a lesser extent. + + Example (optimizing a module with Conv->Batchnorm):: + + import torch + in_channels, out_channels = 3, 32 + conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=True) + bn = torch.nn.BatchNorm2d(out_channels, eps=.001) + mod = torch.nn.Sequential(conv, bn) + frozen_mod = torch.jit.optimize_for_inference(torch.jit.script(mod.eval())) + assert "batch_norm" not in str(frozen_mod.graph) + # if built with MKLDNN, convolution will be run with MKLDNN weights + assert "MKLDNN" in frozen_mod.graph + """ + if not isinstance(mod, ScriptModule): + raise RuntimeError( + "optimize_for_inference expects a ScriptModule as input. " + "Please use torch.jit.script or torch.jit.trace to script your 'nn.Module'." + ) + + if other_methods is None: + other_methods = [] + + if hasattr(mod, "training"): + mod = freeze(mod.eval(), preserved_attrs=other_methods) + + torch._C._jit_pass_optimize_for_inference(mod._c, other_methods) + + return mod diff --git a/lib/python3.12/site-packages/torch/jit/_fuser.py b/lib/python3.12/site-packages/torch/jit/_fuser.py new file mode 100644 index 0000000000000000000000000000000000000000..7466800402d27a954be8f5631df0c9517d10189a --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_fuser.py @@ -0,0 +1,161 @@ +# mypy: allow-untyped-defs +import contextlib +from typing import List, Tuple + +import torch + + +@contextlib.contextmanager +def optimized_execution(should_optimize): + """Context manager that controls whether the JIT's executor will run optimizations before executing a function.""" + stored_flag = torch._C._get_graph_executor_optimize() + torch._C._set_graph_executor_optimize(should_optimize) + try: + yield + finally: + torch._C._set_graph_executor_optimize(stored_flag) + + +@contextlib.contextmanager +def fuser(name): + """Context manager that facilitates switching between backend fusers. + + Valid names: + * ``fuser0`` - enables only legacy fuser + * ``fuser1`` - enables only NNC + * ``fuser2`` - enables only nvFuser + * ``fuser3`` - enables oneDNN Graph + """ + old_cpu_fuse = torch._C._jit_can_fuse_on_cpu() + old_gpu_fuse = torch._C._jit_can_fuse_on_gpu() + old_texpr_fuser_state = torch._C._jit_texpr_fuser_enabled() + old_nvfuser_state = torch._C._jit_nvfuser_enabled() + old_llga_state = torch._C._jit_llga_enabled() + if name == "fuser0": # legacy fuser + torch._C._jit_override_can_fuse_on_cpu(True) + torch._C._jit_override_can_fuse_on_gpu(True) + torch._C._jit_set_texpr_fuser_enabled(False) + torch._C._jit_set_nvfuser_enabled(False) + torch._C._jit_set_llga_enabled(False) + elif name == "fuser1": # NNC + old_profiling_executor = torch._C._jit_set_profiling_executor(True) + old_profiling_mode = torch._C._get_graph_executor_optimize(True) + torch._C._jit_override_can_fuse_on_cpu(True) + torch._C._jit_override_can_fuse_on_gpu(True) + torch._C._jit_set_texpr_fuser_enabled(True) + torch._C._jit_set_nvfuser_enabled(False) + torch._C._jit_set_llga_enabled(False) + elif name == "fuser2": # nvFuser + torch._C._jit_override_can_fuse_on_cpu(False) + torch._C._jit_override_can_fuse_on_gpu(False) + torch._C._jit_set_texpr_fuser_enabled(False) + torch._C._jit_set_nvfuser_enabled(True) + torch._C._jit_set_llga_enabled(False) + elif name == "fuser3": # oneDNN Graph + old_profiling_executor = torch._C._jit_set_profiling_executor(True) + old_profiling_mode = torch._C._get_graph_executor_optimize(True) + torch._C._jit_override_can_fuse_on_cpu(True) + torch._C._jit_override_can_fuse_on_gpu(False) + torch._C._jit_set_texpr_fuser_enabled(True) + torch._C._jit_set_nvfuser_enabled(False) + torch._C._jit_set_llga_enabled(True) + elif name == "none": # Turn Pytorch fuser off + torch._C._jit_override_can_fuse_on_cpu(False) + torch._C._jit_override_can_fuse_on_gpu(False) + torch._C._jit_set_texpr_fuser_enabled(False) + torch._C._jit_set_nvfuser_enabled(False) + torch._C._jit_set_llga_enabled(False) + else: + raise Exception(f"unrecognized fuser option (name: {name})") # noqa: TRY002 + try: + yield + finally: + if name in ["fuser1", "fuser3"]: # NNC or oneDNN Graph + torch._C._jit_set_profiling_executor(old_profiling_executor) # type: ignore[possibly-undefined] + torch._C._get_graph_executor_optimize(old_profiling_mode) # type: ignore[possibly-undefined] + # recover the previous values + torch._C._jit_override_can_fuse_on_cpu(old_cpu_fuse) + torch._C._jit_override_can_fuse_on_gpu(old_gpu_fuse) + torch._C._jit_set_texpr_fuser_enabled(old_texpr_fuser_state) + torch._C._jit_set_nvfuser_enabled(old_nvfuser_state) + torch._C._jit_set_llga_enabled(old_llga_state) + + +last_executed_optimized_graph = torch._C._last_executed_optimized_graph + + +def _get_differentiable_graph_node(node, diff_node): + if node.kind() == "prim::DifferentiableGraph": + diff_node.append(node) + else: + for block in node.blocks(): + for n in block.nodes(): + _get_differentiable_graph_node(n, diff_node) + + +def _graph_for(self, *args, **kwargs): + return _script_method_graph_for(self, self, *args, **kwargs) + + +def _script_method_graph_for(self, parent, *args, **kwargs): + try: + dbs = parent.get_debug_state() + eps = list(dbs.execution_plans.values()) + assert len(eps) == 1 + graph = eps[0].graph.copy() + + # graph_executor_states for differentiable node + fw_states = eps[0].code.differentiable_op_executor_states() + diff_nodes: List[torch._C.Node] = [] + for n in graph.nodes(): + _get_differentiable_graph_node(n, diff_nodes) + + assert len(fw_states) == len(diff_nodes) + # swap each differentiable graph with optimized graph in their execution plan + for n, state in zip(diff_nodes, fw_states): + fw_execution_plans = list(state.execution_plans.values()) + # we can only update the subgraph when there's a unique execution + # plan. Avoid assert here so we would skip the ones that can't be + # updated while try the best effort to update other nodes. + if len(fw_execution_plans) == 1: + n.g_("Subgraph", fw_execution_plans[0].graph) + + return graph + except Exception: + # fallback approach, we just ran the graph and return the recorded optimized + # graph + self(*args, **kwargs) + return last_executed_optimized_graph() + + +def set_fusion_strategy(strategy: List[Tuple[str, int]]): + """Set the type and number of specializations that can occur during fusion. + + Usage: provide a list of pairs (type, depth) where type is one of "STATIC" or "DYNAMIC" + and depth is an integer. + + Behavior - static vs dynamic: + In STATIC fusion, fused ops are compiled to have fixed input shapes. The shape is determined + based on some initial profiling runs. + In DYNAMIC fusion, fused ops are compiled to have variable input shapes, so that multiple + shapes are possible. + + In both cases, we also recompile on new striding behavior, device, or dtype. + + Behavior - fallback functions & depth: + When an input doesn't match the format required by the specialized compiled op, it will run + a fallback function. Fallback functions are recursively be compiled and specialized based + on the observed tensor shapes. Since compilation can be slow, the "depth" parameter is provided to + limit the number of specializations that can be compiled, before giving up on recompiling and + falling back to a completely un-fused, un-specialized implementation. + + The list of (type, depth) pairs controls the type of specializations and the number of + specializations. For example: [("STATIC", 2), ("DYNAMIC", 2)] indicates that the first + two specializations will use static fusions, the following two specializations will use + dynamic fusion, and any inputs that satisfy none of the 4 options will run an + unfused implementation. + + NB: in the future, if more as more fusion backends are added there may be more granular + apis for specific fusers. + """ + return torch._C._jit_set_fusion_strategy(strategy) diff --git a/lib/python3.12/site-packages/torch/jit/_ir_utils.py b/lib/python3.12/site-packages/torch/jit/_ir_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..52b953624a3a18d2fb3f06f79044ab6ac583af76 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_ir_utils.py @@ -0,0 +1,26 @@ +# mypy: allow-untyped-defs +from typing import Union + +import torch + + +class _InsertPoint: + def __init__( + self, + insert_point_graph: torch._C.Graph, + insert_point: Union[torch._C.Node, torch._C.Block], + ): + self.insert_point = insert_point + self.g = insert_point_graph + self.guard = None + + def __enter__(self): + self.prev_insert_point = self.g.insertPoint() + self.g.setInsertPoint(self.insert_point) + + def __exit__(self, *args): + self.g.setInsertPoint(self.prev_insert_point) + + +def insert_point_guard(self, insert_point: Union[torch._C.Node, torch._C.Block]): + return _InsertPoint(self, insert_point) diff --git a/lib/python3.12/site-packages/torch/jit/_passes/__init__.py b/lib/python3.12/site-packages/torch/jit/_passes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/jit/_passes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/_passes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c0a98a9ce47b438d3d223ec44c3ae68d808bdd2a Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/_passes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/_passes/__pycache__/_property_propagation.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/_passes/__pycache__/_property_propagation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..554e5cd75f54f628a154389f4717fc769231c1b7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/_passes/__pycache__/_property_propagation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/_passes/_property_propagation.py b/lib/python3.12/site-packages/torch/jit/_passes/_property_propagation.py new file mode 100644 index 0000000000000000000000000000000000000000..1537f7bc4147359af3f1e0ea4229352d37303921 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_passes/_property_propagation.py @@ -0,0 +1,47 @@ +# mypy: allow-untyped-defs +""" +Tools to help with tensor property propagation. + +This is not intended to be imported directly; please use the exposed +functionalities in `torch.jit`. +""" + +from typing import Any, List + +import torch +from torch import TensorType +from torch._C import Graph + + +def apply_input_props_using_example(graph: Graph, example_input: List[Any]): + """ + Applies properties for each tensor in the graph inputs + using the example supplied. + """ + graph_inputs = list(graph.inputs()) + if len(graph_inputs) == 0: + return + + # Strip self args off for methods + in_0 = graph_inputs[0] + if isinstance(in_0.type(), torch._C.ClassType) and in_0.debugName() == "self": + graph_inputs = graph_inputs[1:] + + if not len(graph_inputs) == len(example_input): + raise RuntimeError( + "Number of inputs in graph does not match number of inputs in the example" + ) + + for i, (graph_i, example_i) in enumerate(zip(graph_inputs, example_input)): + if example_i is None: + continue # Skip the type check + + if isinstance(example_i, torch.Tensor) != isinstance( + graph_i.type(), TensorType + ): + raise RuntimeError( + f"Input {i} does not match type of example", graph_i, example_i + ) + + if isinstance(example_i, torch.Tensor): + graph_i.setType(TensorType.create_from_tensor(example_i)) # type: ignore[arg-type] diff --git a/lib/python3.12/site-packages/torch/jit/_pickle.py b/lib/python3.12/site-packages/torch/jit/_pickle.py new file mode 100644 index 0000000000000000000000000000000000000000..5517499e926001603cdb8f0696a7dfdc7df1a964 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_pickle.py @@ -0,0 +1,38 @@ +# mypy: allow-untyped-defs +# These functions are referenced from the pickle archives produced by +# ScriptModule.save() + + +# These (`build_*`) functions used to be used by `pickler.cpp` to specify +# the type of the list for certain special types, but now all lists get +# a type attached and restored via `restore_type_tag` below. The legacy +# functions should stick around for backwards-compatibility. + + +def build_intlist(data): + return data + + +def build_tensorlist(data): + return data + + +def build_doublelist(data): + return data + + +def build_boollist(data): + return data + + +def build_tensor_from_id(data): + if isinstance(data, int): + # just the id, can't really do anything + return data + + +def restore_type_tag(value, type_str): + # The type_ptr is used by the jit unpickler to restore the full static type + # to container types like list when they are re-loaded, but this doesn't + # matter for Python, so just return the plain value + return value diff --git a/lib/python3.12/site-packages/torch/jit/_script.py b/lib/python3.12/site-packages/torch/jit/_script.py new file mode 100644 index 0000000000000000000000000000000000000000..b572beca7d86db3110bb31be1a77a5b17454cc74 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_script.py @@ -0,0 +1,1726 @@ +"""TorchScript. + +This module contains functionality to support the JIT's scripting frontend, notably: + - torch.jit.script + +This is not intended to be imported directly; please use the exposed +functionalities in `torch.jit`. +""" +import collections +import copy +import enum +import functools +import inspect +import pickle +import warnings +from typing import Any, Callable, Dict, List, Set, Tuple, Union + +import torch +import torch._jit_internal as _jit_internal +from torch._classes import classes +from torch._jit_internal import _get_model_id, _qualified_name +from torch._utils_internal import log_torchscript_usage +from torch.jit._builtins import _register_builtin +from torch.jit._fuser import _graph_for, _script_method_graph_for +from torch.jit._monkeytype_config import ( + JitTypeTraceConfig, + JitTypeTraceStore, + monkeytype_trace, +) +from torch.jit._recursive import ( + _compile_and_register_class, + infer_methods_to_compile, + ScriptMethodStub, + wrap_cpp_module, +) +from torch.jit._state import ( + _enabled, + _set_jit_function_cache, + _set_jit_overload_cache, + _try_get_jit_cached_function, + _try_get_jit_cached_overloads, +) +from torch.jit.frontend import get_default_args, get_jit_class_def, get_jit_def +from torch.nn import Module +from torch.overrides import ( + has_torch_function, + has_torch_function_unary, + has_torch_function_variadic, +) +from torch.package import PackageExporter, PackageImporter +from torch.utils import set_module + +from ._serialization import validate_map_location + + +type_trace_db = JitTypeTraceStore() # DB to hold all call traces from MonkeyType + +torch._C.ScriptMethod.graph_for = _script_method_graph_for # type: ignore[attr-defined] +torch._C.ScriptFunction.graph_for = _graph_for # type: ignore[attr-defined] +ScriptFunction = torch._C.ScriptFunction +ScriptFunction.__doc__ = """ +Functionally equivalent to a :class:`ScriptModule`, but represents a single +function and does not have any attributes or Parameters. +""" +set_module(ScriptFunction, "torch.jit") + + +# Throws an error if a jit function is pickled. +# Helps to avoid Python crashes for Python versions 3.9.5 + when protocol 0 or 1 is given as an argument. +def _reduce(cls): + raise pickle.PickleError("ScriptFunction cannot be pickled") + + +ScriptFunction.__reduce__ = _reduce # type: ignore[assignment] + + +if _enabled: + Attribute = collections.namedtuple("Attribute", ["value", "type"]) +else: + + def Attribute(value, type): # type: ignore[no-redef] + return value + + +Attribute.__doc__ = """ + This method is a pass-through function that returns `value`, mostly + used to indicate to the TorchScript compiler that the left-hand side + expression is a class instance attribute with type of `type`. Note that + `torch.jit.Attribute` should only be used in `__init__` method of `jit.ScriptModule` + subclasses. + + Though TorchScript can infer correct type for most Python expressions, there are some cases where + type inference can be wrong, including: + + - Empty containers like `[]` and `{}`, which TorchScript assumes to be container of `Tensor` + - Optional types like `Optional[T]` but assigned a valid value of type `T`, TorchScript would assume + it is type `T` rather than `Optional[T]` + + In eager mode, it is simply a pass-through function that returns `value` + without other implications. + + Example: + + .. testcode:: + + import torch + from typing import Dict + + class AttributeModule(torch.jit.ScriptModule): + def __init__(self) -> None: + super().__init__() + self.foo = torch.jit.Attribute(0.1, float) + + # we should be able to use self.foo as a float here + assert 0.0 < self.foo + + self.names_ages = torch.jit.Attribute({}, Dict[str, int]) + self.names_ages["someone"] = 20 + assert isinstance(self.names_ages["someone"], int) + + m = AttributeModule() + # m will contain two attributes + # 1. foo of type float + # 2. names_ages of type Dict[str, int] + + .. testcleanup:: + + del AttributeModule + del m + + Note: it's now preferred to instead use type annotations instead of `torch.jit.Attribute`: + + .. testcode:: + + import torch + from typing import Dict + + class AttributeModule(torch.nn.Module): + names: Dict[str, int] + + def __init__(self) -> None: + super().__init__() + self.names = {} + + m = AttributeModule() + + .. testcleanup:: + + del AttributeModule + del m + + Args: + value: An initial value to be assigned to attribute. + type: A Python type + + Returns: + Returns `value` +""" + + +def _get_type_trace_db(): + # This is a private API. Use of this for external purposes is discouraged. + return type_trace_db + + +# Gets a function from the name of a method on a type +def _get_function_from_type(cls, name): + return getattr(cls, name, None) + + +# ScriptClasses must be new-style classes because we construct them using their +# __new__ method. +def _is_new_style_class(cls): + if hasattr(cls, "__class__"): + return "__dict__" in dir(cls) or hasattr(cls, "__slots__") + + +# These OrderedDictWrapper classes replace the actual OrderedDicts in +# module with versions that get/set properties inside of Module. +# This allows us to reuse most of nn.Module while still storing the +# data in C++. +# Each OrderedDict needs to support: +# x not in view +# x in view +# view[name] = ... +# view.values() +# del view[name] +# view.items() +# view.keys() +# len(view) + + +class OrderedDictWrapper: + def __init__(self, _c): + self._c = _c + + def keys(self): + return [k for k, v in self.items()] + + def values(self): + return [v for k, v in self.items()] + + def __len__(self): + return len(self.values()) + + def __delitem__(self, k): + raise RuntimeError("cannot delete methods or parameters of a script module") + + def items(self): + return self._c.items() + + def __setitem__(self, k, v): + if k not in self: + raise RuntimeError( + f"Can't add a new parameter after ScriptModule construction. Tried to add '{k}" + ) + self._c.setattr(k, v) + + def __contains__(self, k): + return self._c.contains(k) + + def __getitem__(self, k): + if k not in self: + raise KeyError(k) + return self._c.getattr(k) + + +class OrderedModuleDict(OrderedDictWrapper): + def __init__(self, module, python_dict): + super().__init__(torch._C.ModuleDict(module)) + # contains _both_ script modules and non-script python-only modules + + # because script modules are subclassed in python and the + # C++ Module class will not hold references to them, + # to ensure that you always get the same python value here + # we store it in the python dict as well + self._python_modules = python_dict + + def items(self): + r = self._python_modules.items() + return r + + def __contains__(self, k): + return k in self._python_modules + + def __setitem__(self, k, v): + # Cases where sub-module can be re-assigned after ScriptModule construction + # 1. If the attr is an module interface type, it's guaranteed that the module is + # not inlined in the graph, so it's safe to swap a new ScriptModule in. + # 2. if the new value if a ScriptModule with the same JIT type, IR won't change + # and it's legit to swap a new module in. + # In these two cases we allow swapping a new scripted module and update the + # corresponding python module dict to keep sync. + # Note: the value to be swapped in has to be ScriptModule instead of nn.Module, + # otherwise it's illegal and we throw error. + if isinstance(v, ScriptModule): + self._c.setattr(k, v) + self._python_modules[k] = v + else: + raise RuntimeError( + "Cannot re-assign modules in a ScriptModule with non-scripted " + f"module, tried to replace existing module '{k}': {v}" + ) + + def __getitem__(self, k): + return self._python_modules[k] + + +# For each user-defined class that subclasses ScriptModule, this meta-class: +# (1) finds all the methods annotated with @script_method in a ScriptModule and +# removes them from the class attributes +# (2) puts a wrapper around the class's __init__ method to recursively compile +# all of the script_methods with the module after the original __init__ has +# run. This has to occur after the user-defined __init__ so that submodules and +# parameters are initialized _before_ the script compiler resolve references to +# `self.param` or `self.module`. +class ScriptMeta(type): + def __init__(cls, name, bases, attrs): # noqa: B902 + # Aggregate all the ScriptMethods and constants from superclasses + cls._methods: Dict[str, Any] = {} + cls._constants_set = set(getattr(cls, "__constants__", ())) + for base in reversed(bases): + for k, v in getattr(base, "_methods", {}).items(): + cls._methods[k] = v + base_constants: Set = getattr(base, "_constants_set", set()) + cls._constants_set = cls._constants_set.union(base_constants) + + # find all the script methods of the current class + for k, v in sorted(attrs.items()): + if isinstance(v, ScriptMethodStub): + delattr(cls, k) + cls._methods[v.original_method.__name__] = v + + if getattr(cls, "_disable_script_meta", False): + # We leave built-in ScriptModule types alone, since this metaclass + # is only for compiling user classes that inherit from + # ScriptModule. + super().__init__(name, bases, attrs) + return + + original_init = getattr(cls, "__init__", lambda self: None) + + @functools.wraps(original_init) + def init_then_script(self, *args, **kwargs): + num_methods = len(cls._methods) + original_init(self, *args, **kwargs) + added_methods_in_init = len(cls._methods) > num_methods + + if type(self) == cls: + + def make_stubs(module): + cls = type(module) + if hasattr(cls, "_methods"): + return [v for k, v in sorted(cls._methods.items())] + else: + return infer_methods_to_compile(module) + + self.__dict__[ + "_actual_script_module" + ] = torch.jit._recursive.create_script_module( + self, make_stubs, share_types=not added_methods_in_init + ) + + # Delete the Python attributes that now shadow the ScriptModule + # ones, so that __getattr__ and __setattr__ will properly find + # the scripted versions. + concrete_type = self._actual_script_module._concrete_type + for name in concrete_type.get_attributes(): + delattr(self, name) + for name, _ in concrete_type.get_modules(): + delattr(self, name) + for name in ("_parameters", "_buffers", "_modules"): + delattr(self, name) + + cls.__init__ = init_then_script # type: ignore[misc] + super().__init__(name, bases, attrs) + + +class _CachedForward: + def __get__(self, obj, cls): + return self.__getattr__("forward") # type: ignore[attr-defined] + + +class ScriptWarning(Warning): + pass + + +def script_method(fn): + if not _enabled: + return fn + # NOTE: we need to traverse two frames here because the meta-class frame + # for ScriptModule will be present, as opposed to invoking @script on a + # a function or invoking define() on a CompilationUnit. + # The stack will look like: + # + # 0. createResolutionCallback() + # 1. script_method() + # 2. ScriptModule metaclass frame + # 3. Surrounding scope + # + # createResolutionCallback internally adds 1 to get us to the scope of this + # function (the calling function). Adding 2 gets us to the proper surrounding scope. + _rcb = _jit_internal.createResolutionCallbackFromFrame(frames_up=2) + ast = get_jit_def(fn, fn.__name__, self_name="ScriptModule") + return ScriptMethodStub(_rcb, ast, fn) + + +class ConstMap: + def __init__(self, const_mapping): + self.const_mapping = const_mapping + + def __getattr__(self, attr): + return self.const_mapping[attr] + + +def unpackage_script_module( + importer: PackageImporter, script_module_id: str +) -> torch.nn.Module: + """ + Call by ``torch.package.PackageImporter``'s Pickler's ``persistent_load`` function. + + Performs work of loading and returning a ScriptModule from a ``torch.package`` archive. + """ + if not isinstance(importer.zip_reader, torch._C.PyTorchFileReader): + raise RuntimeError( + "Loading ScriptObjects from a PackageImporter created from a " + "directory is not supported. Use a package archive file instead." + ) + cu = torch._C.CompilationUnit() + cpp_module = torch._C._import_ir_module_from_package( + cu, + importer.zip_reader, + importer.storage_context, + validate_map_location(importer.last_map_location), + script_module_id, + ) + return wrap_cpp_module(cpp_module) + + +if _enabled: + _magic_methods = [ + "__iter__", + "__len__", + "__neg__", + "__mul__", + "__contains__", + "__add__", + "__sub__", + "__pow__", + "__truediv__", + "__mod__", + "__ne__", + "__eq__", + "__lt__", + "__gt__", + "__le__", + "__ge__", + "__and__", + "__or__", + "__xor__", + "__getitem__", + "__setitem__", + "__call__", + "__int__", + "__float__", + "__bool__", + "__str__", + "__enter__", + "__exit__", + ] + + class RecursiveScriptClass: + """Wrapper for a TorchScript class instance for use in Python. + + An analogue of RecursiveScriptModule for regular objects that are not modules. + This class is a wrapper around a torch._C.ScriptObject that represents an instance + of a TorchScript class and allows it to be used in Python. + + Attributes: + _c [torch._C.ScriptObject]: The C++ object to which attribute lookups and method + calls are forwarded. + _props [Dict[str, property]]: A dictionary of properties fetched from self._c and + exposed on this wrppaer. + """ + + def __init__(self, cpp_class): + super().__init__() + self.__dict__["_initializing"] = True + self._c = cpp_class + + # Add wrapped object's properties to this class instance. + self._props = { + prop.name: property(prop.getter, prop.setter) + for prop in self._c._properties() + } + + self.__dict__["_initializing"] = False + + def __getattr__(self, attr): + if self.__dict__.get("_initializing"): + return super().__getattr__(attr) # type: ignore[misc] + + if attr in self._props: + return self._props[attr].fget() # type: ignore[call-arg, misc] + + return getattr(self._c, attr) + + def __setattr__(self, attr, value): + if self.__dict__.get("_initializing"): + return super().__setattr__(attr, value) + + if attr in self._props: + return self._props[attr].fset(value) # type: ignore[call-arg, misc] + + setattr(self._c, attr, value) + + # Delegate calls to magic methods like __len__ to the C++ module backing the + # RecursiveScriptClass. + def forward_magic_method(self, method_name, *args, **kwargs): + if not self._c._has_method(method_name): + raise TypeError + + self_method = self.__getattr__(method_name) + return self_method(*args, **kwargs) + + def __getstate__(self): + raise pickle.PickleError("ScriptClasses cannot be pickled") + + def __iadd__(self, other): + if self._c._has_method("__iadd__"): + return self.forward_magic_method("__iadd__", other) + else: + return self.forward_magic_method("__add__", other) + + for method_name in _magic_methods: + + def method_template(self, *args, **kwargs): + return self.forward_magic_method(method_name, *args, **kwargs) + + setattr(RecursiveScriptClass, method_name, method_template) + + # this is a Python 'non-data descriptor' that causes the first access + # to ScriptModule's forward to look up the forward method and stash + # it in the objects dict. Due to the standard rules for attribute lookup, + # subsequent lookups will just directly return the previously looked up method. + # This is necessary because nn.Module defines forward as a method. If we + # did nothing, __getattr__ would not be called. Instead we'd get nn.Module.forward + # which always throws an exception. + + class ScriptModule(Module, metaclass=ScriptMeta): + r"""Wrapper for C++ torch::jit::Module with methods, attributes, and parameters. + + A wrapper around C++ ``torch::jit::Module``. ``ScriptModule``\s + contain methods, attributes, parameters, and + constants. These can be accessed the same way as on a normal ``nn.Module``. + """ + + __jit_unused_properties__ = [ + "code", + "code_with_constants", + "graph", + "inlined_graph", + "original_name", + ] + + def __init__(self) -> None: + super().__init__() + + forward: Callable[..., Any] = _CachedForward() # type: ignore[assignment] + + def __getattr__(self, attr): + if "_actual_script_module" not in self.__dict__: + return super().__getattr__(attr) + return getattr(self._actual_script_module, attr) + + def __setattr__(self, attr, value): + if "_actual_script_module" not in self.__dict__: + # Unwrap torch.jit.Attribute into a regular setattr + record + # the provided type in __annotations__. + # + # This ensures that if we use the attr again in `__init__`, it + # will look like the actual value, not an instance of Attribute. + if isinstance(value, Attribute): + # NB: Ensure that we set __annotations__ on the specific + # class in question, and not on a superclass (which would + # be wrong wrong wrong!). + # See also https://github.com/pytorch/pytorch/issues/39463 + if "__annotations__" not in self.__class__.__dict__: + self.__class__.__annotations__ = {} + self.__annotations__[attr] = value.type + value = value.value + return super().__setattr__(attr, value) + + setattr(self._actual_script_module, attr, value) + + def define(self, src): + if "_actual_script_module" in self.__dict__: + # If we have completed initialization, just defer to the + # backing RecursiveScriptModule to eagerly compile the provided + # source. + return self._actual_script_module.define(src) + + # Otherwise, we are still in the object's __init__. + # In that case, add `src` as a stub to be compiled. + # + # We use frames_up=1 to get to the proper surrounding scope. The stack + # will look like: + # 0. createResolutionCallback + # 1. define() + # 2. surrounding scope. + # + # createResolutionCallback internally adds 1 to get us to our frame, then + # we add 1 to get to the proper surrounding scope. + rcb = _jit_internal.createResolutionCallbackFromFrame(frames_up=1) + ast = torch._C._parse_source_def(src) + self._methods[ast.name().name] = ScriptMethodStub(rcb, ast, None) + + def _replicate_for_data_parallel(self): + return self._actual_script_module._replicate_for_data_parallel() + + def __reduce_package__(self, exporter: PackageExporter): + """Save a ScriptModule inside of a ``torch.package`` archive. + + Called by ``torch.package.PackageExporter``'s Pickler's ``persistent_id`` when + saving TorchScript objects. Performs act of saving a ScriptModule inside of + a ``torch.package`` archive. + + Returns method to load the ScriptModule from a ``torch.package.PackageImporter``'s + Pickler's ``persistent_load`` function. + """ + script_module_id = exporter.get_unique_id() + exporter.script_module_serializer.serialize(self._c, int(script_module_id)) + return (unpackage_script_module, (script_module_id,)) + + class RecursiveScriptModule(ScriptModule): + # XXX: RecursiveScriptModule inherits from ScriptModule for the sole + # reason that it retains the existing isinstance(ScriptModule) + # behavior. + r"""Retain the existing isinstance(ScriptModule) behavior. + + The core data structure in TorchScript is the ``ScriptModule``. It is an + analogue of torch's ``nn.Module`` and represents an entire model as a tree of + submodules. Like normal modules, each individual module in a ``ScriptModule`` can + have submodules, parameters, and methods. In ``nn.Module``\s methods are implemented + as Python functions, but in ``ScriptModule``\s methods are implemented as + TorchScript functions, a statically-typed subset of Python that contains all + of PyTorch's built-in Tensor operations. This difference allows your + ``ScriptModule``\s code to run without the need for a Python interpreter. + + ``ScriptModule``\s should not be created manually, instead use + either :func:`tracing ` or :func:`scripting `. + Tracing and scripting can be applied incrementally and :ref:`composed as necessary `. + + * Tracing records the tensor operations as executed with a set of example inputs and uses these + operations to construct a computation graph. You can use the full dynamic behavior of Python with tracing, + but values other than Tensors and control flow aren't captured in the graph. + + * Scripting inspects the Python code of the model + and compiles it to TorchScript. Scripting allows the use of many `types`_ of values and supports dynamic control flow. + Many, but not all features of Python are supported by the compiler, so changes to the source code may be necessary. + """ + + _disable_script_meta = True + + def __init__(self, cpp_module): + self.__dict__["_initializing"] = True + self._c = cpp_module + super().__init__() + # Delete the 'training' attribute set up by `Module.__init__`. It + # will get set on the underlying cpp module, so we delete it here + # to avoid this version shadowing the cpp module version. + delattr(self, "training") + + @staticmethod + def _construct(cpp_module, init_fn): + """ + Construct a RecursiveScriptModule that's ready for use. + + PyTorch code should use this to construct a RecursiveScriptModule instead + of instead of calling `__init__` directly, as it makes sure the + object is properly finalized (and in the future, we may take + control of how the RecursiveScriptModule instance is created). + + Args: + cpp_module: The C++ Module that will hold the actual state of + this RecursiveScriptModule instance. + init_fn: Lambda that initializes the RecursiveScriptModule passed to it. + """ + script_module = RecursiveScriptModule(cpp_module) + init_fn(script_module) + + # Finalize the ScriptModule: replace the nn.Module state with our + # custom implementations and flip the _initializing bit. + RecursiveScriptModule._finalize_scriptmodule(script_module) + return script_module + + @staticmethod + def _finalize_scriptmodule(script_module): + script_module._parameters = OrderedDictWrapper( + torch._C.ParameterDict(script_module._c) + ) + script_module._buffers = OrderedDictWrapper( + torch._C.BufferDict(script_module._c) + ) + script_module._modules = OrderedModuleDict( + script_module._c, script_module._modules + ) + script_module._initializing = False + + def _reconstruct(self, cpp_module): + """ + Re-construct an instance of RecursiveScriptModule using an instance of a C++ module. + + Args: + cpp_module: The C++ module that this RecursiveScriptModule will be rebuilt around. + """ + self.__init__(cpp_module) # type: ignore[misc] + + # Copy the concrete type from the C++ module to this ScriptModule. + self._concrete_type = torch._C.ConcreteModuleType.from_jit_type( + self._c._type() + ) + + # Copy submodules from the C++ module to this ScriptModule. + modules = {} + for name, cpp_module in torch._C.ModuleDict(self._c).items(): + modules[name] = wrap_cpp_module(cpp_module) + self._modules = OrderedModuleDict(self._c, modules) # type: ignore[assignment] + + # Copy parameters and buffers. + self._parameters = OrderedDictWrapper(torch._C.ParameterDict(self._c)) # type: ignore[assignment] + self._buffers = OrderedDictWrapper(torch._C.BufferDict(self._c)) # type: ignore[assignment] + + # Get rid of the functions from the old C++ module. + self.__dict__ = { + k: v + for k, v in self.__dict__.items() + if not isinstance(v, torch._C.ScriptMethod) + } + self.__dict__["_initializing"] = False + + @property + def graph(self): + r"""Return a string representation of the internal graph for the ``forward`` method. + + See :ref:`interpreting-graphs` for details. + """ + return self._c._get_method("forward").graph + + @property + def inlined_graph(self): + r""" + Return a string representation of the internal graph for the ``forward`` method. + + This graph will be preprocessed to inline all function and method calls. + See :ref:`interpreting-graphs` for details. + """ + return self.forward.inlined_graph # type: ignore[attr-defined] + + @property + def code(self): + r""" + Return a pretty-printed representation (as valid Python syntax) of the internal graph for the ``forward`` method. + + See :ref:`inspecting-code` for details. + """ + return self.forward.code # type: ignore[attr-defined] + + @property + def code_with_constants(self): + r"""Return a tuple. + + Returns a tuple of: + + [0] a pretty-printed representation (as valid Python syntax) of + the internal graph for the ``forward`` method. See `code`. + [1] a ConstMap following the CONSTANT.cN format of the output in [0]. + The indices in the [0] output are keys to the underlying constant's values. + + See :ref:`inspecting-code` for details. + """ + r = self.forward.code_with_constants # type: ignore[attr-defined] + return (r[0], ConstMap(r[1])) + + def save(self, f, **kwargs): + r"""Save with a file-like object. + + save(f, _extra_files={}) + + See :func:`torch.jit.save ` which accepts a file-like object. + This function, torch.save(), converts the object to a string, treating it as a path. + DO NOT confuse these two functions when it comes to the 'f' parameter functionality. + """ + return self._c.save(str(f), **kwargs) + + def _save_for_lite_interpreter(self, *args, **kwargs): + r"""Add (or update) the bytecode session to the script model. + + _save_for_lite_interpreter(f) + + The updated model is used + in lite interpreter for mobile applications. + + Args: + f: a string containing a file name. + _extra_files: Map from filename to contents which will be stored as part of 'f'. + + """ + return self._c._save_for_mobile(*args, **kwargs) + + def _save_to_buffer_for_lite_interpreter(self, *args, **kwargs): + return self._c._save_to_buffer_for_mobile(*args, **kwargs) + + def save_to_buffer(self, *args, **kwargs): + return self._c.save_to_buffer(*args, **kwargs) + + def get_debug_state(self, *args, **kwargs): + return self._c.get_debug_state() + + def extra_repr(self): + return f"original_name={self.original_name}" + + def graph_for(self, *args, **kwargs): + return self.forward.graph_for(self, *args, **kwargs) # type: ignore[attr-defined] + + @property + def original_name(self): + if type(self) == str(self._c._type().name()): + return "" + return str(self._c._type().name()) + + def define(self, src): + # We use frames_up=1 to get to the proper surrounding scope. The stack + # will look like: + # 0. createResolutionCallback + # 1. define() + # 2. surrounding scope. + # + # createResolutionCallback internally adds 1 to get us to our frame, then + # we add 1 to get to the proper surrounding scope. + rcb = _jit_internal.createResolutionCallbackFromFrame(frames_up=1) + self._c._define(self._concrete_type, src, rcb) + + def __getattr__(self, attr): + if "_initializing" not in self.__dict__: + raise RuntimeError( + "ScriptModule has not been initialized, did you forget to call super's init?" + ) + + if self._initializing: + return super().__getattr__(attr) + + # _modules check is before hasattr since modules are included as attributes in _c, + # but we want to get the python wrapper from _modules instead of the raw _c object. + if attr in self._modules: + return self._modules[attr] + elif self._c.hasattr(attr): + return self._c.getattr(attr) + elif self._c._has_method(attr): + script_method = self._c._get_method(attr) + # cache method so future calls do not go through __getattr__ + # to improve invocation performance + self.__dict__[attr] = script_method + return script_method + + return super().__getattr__(attr) + + def __setattr__(self, attr, value): + if self._initializing: + return super().__setattr__(attr, value) + + if attr in self._modules: + self._modules[attr] = value + elif self._c.hasattr(attr): + self._c.setattr(attr, value) + elif ( + hasattr(self, "_concrete_type") + and attr in self._concrete_type.get_constants().keys() + ): + # TODO: we don't have _concrete_type set after load(), and in general we lose constant information. + # We should encode constants as class type attributes (or something) so it persists across save/load. + raise AttributeError( + f"Cannot mutate TorchScript constant value: '{attr}'. Value: '{value}'" + ) + else: + # We allow setting Python attributes on the ScriptModule, for + # when people want to stash some convenience info on it. + # TODO: it's possible that the following is confusing: + # s = torch.jit.script(...) + # s.python_attr = ... + # s.save() <--- this doesn't have `python_attr` + # It's fairly trivial to save enough info to warn in this case. + return super().__setattr__(attr, value) + + def __copy__(self): + return torch.jit._recursive.wrap_cpp_module(copy.copy(self._c)) + + def __deepcopy__(self, memo): + return torch.jit._recursive.wrap_cpp_module(copy.deepcopy(self._c, memo)) + + # Python magic methods do method lookups on an object's class type, instead of looking up + # the method defines on the class instance. In order to continue to expose the magic methods + # of builtin-containers (ModuleList, Sequential, ModuleDict) to Python, we + # define magic methods here as a shim to the correct attribute. + def forward_magic_method(self, method_name, *args, **kwargs): + self_method = getattr(self, method_name) + if getattr(self_method, "__func__", None) == getattr( + RecursiveScriptModule, method_name + ): + raise NotImplementedError + return self_method(*args, **kwargs) + + def __iter__(self): + return self.forward_magic_method("__iter__") + + def __getitem__(self, idx): + return self.forward_magic_method("__getitem__", idx) + + def __len__(self): + return self.forward_magic_method("__len__") + + def __contains__(self, key): + return self.forward_magic_method("__contains__", key) + + # dir is defined by the base nn.Module, so instead of throwing if + # it is not overridden, we call into the nn.Module __dir__ method + def __dir__(self): + self_method = self.__dir__ + if ( + self_method.__func__ # type: ignore[attr-defined] + == _get_function_from_type(RecursiveScriptModule, "__dir__") + ): + return super().__dir__() + return self_method() + + # to resolve bool(value), Python looks if __bool__ is defined then __iter__ + # is defined then returns true for classes. Since __iter__() on this + # class throws if it isn't overridden, we define __bool__ to preserve default behavior + def __bool__(self): + self_method = self.__bool__ + if ( + self_method.__func__ # type: ignore[attr-defined] + == _get_function_from_type(RecursiveScriptModule, "__bool__") + ): + return True + return self_method() + + def _replicate_for_data_parallel(self): + # we have to initialize ScriptModule properly so that + # it works with pybind11 + def init_fn(script_module): + # Don't do anything here, we'll initialize the ScriptModule below + return + + return RecursiveScriptModule._construct( + self._c._replicate_for_data_parallel(), init_fn + ) + + # Need to copy all RecursiveScriptModule methods to ScriptModule. + # + # This is because `super().foo()` does not use + # `__getattr__` to look up `foo`. So we need to make each method available on + # the ScriptModule manually. + for name, item in RecursiveScriptModule.__dict__.items(): + if not callable(item) and not isinstance(item, property): + continue + if name.startswith("__") or hasattr(ScriptModule, name): + continue + # We can copy over the implementation wholesale because besides the + # `super()` thing above, ScriptModule behaves exactly like + # RecursiveScriptModule + setattr(ScriptModule, name, item) + + def _get_methods(cls): + import inspect + + # In Python 3 unbound methods are functions, but in Python 2 they are methods + return inspect.getmembers( + cls, predicate=lambda x: inspect.isfunction(x) or inspect.ismethod(x) + ) + + _compiled_methods_allowlist = { + "forward", + "register_buffer", + "register_parameter", + "register_module", + "add_module", + "_apply", + "apply", + "cuda", + "cpu", + "to", + "type", + "float", + "double", + "half", + "state_dict", + "_save_to_state_dict", + "load_state_dict", + "_load_from_state_dict", + "_named_members", + "parameters", + "named_parameters", + "buffers", + "named_buffers", + "children", + "named_children", + "modules", + "named_modules", + "zero_grad", + "share_memory", + "_get_name", + "extra_repr", + "_slow_forward", + "_tracing_name", + "eval", + "train", + "get_extra_state", + "set_extra_state", + } + + def _make_fail(name): + def fail(self, *args, **kwargs): + raise RuntimeError(name + " is not supported on ScriptModules") + + return fail + + for name, method in _get_methods(torch.nn.Module): + if name.startswith("__") or name.endswith("_call_impl"): + continue + if ( + name not in RecursiveScriptModule.__dict__ + and name not in _compiled_methods_allowlist + ): + setattr(RecursiveScriptModule, method.__name__, _make_fail(name)) + + +else: + # TODO MAKE SURE THAT DISABLING WORKS + class RecursiveScriptClass: # type: ignore[no-redef] + pass + + class ScriptModule(torch.nn.Module): # type: ignore[no-redef] + def __init__(self, arg=None): + super().__init__() + + class RecursiveScriptModule(ScriptModule): # type: ignore[no-redef] + def __init__(self, arg=None): + super().__init__() + + +def call_prepare_scriptable_func_impl(obj, memo): + if not isinstance(obj, torch.nn.Module): + return obj + + obj_id = id(obj) + + # If obj_id is in memo, obj has already been prepared or is being + # prepared in another call up the stack. + if obj_id in memo: + return memo[id(obj)] + + obj = obj.__prepare_scriptable__() if hasattr(obj, "__prepare_scriptable__") else obj # type: ignore[operator] + # Record obj in memo to avoid infinite recursion in the case of cycles in the module + # hierarchy when recursing below. + memo[obj_id] = obj + + new_obj_dict = {} + + for name, sub_module in obj.__dict__.items(): + if name == "_modules": + for k, v in sub_module.items(): + sub_module[k] = call_prepare_scriptable_func_impl(v, memo) + new_obj_dict[name] = sub_module + elif isinstance(sub_module, torch.nn.Module) and not isinstance( + sub_module, ScriptModule + ): + new_obj_dict[name] = call_prepare_scriptable_func_impl(sub_module, memo) + else: + new_obj_dict[name] = sub_module + + for k, v in new_obj_dict.items(): + obj.__dict__[name] = v + + return obj + + +def call_prepare_scriptable_func(obj): + memo: Dict[int, torch.nn.Module] = {} + return call_prepare_scriptable_func_impl(obj, memo) + + +def create_script_dict(obj): + """ + Create a ``torch._C.ScriptDict`` instance with the data from ``obj``. + + Args: + obj (dict): The Python dictionary that is used to initialize the ``ScriptDict`` + returned by this function. + + Returns: + An instance of ``torch._C.ScriptDict`` that has the same data as ``obj`` + and can be passed between Python and TorchScript with reference semantics and + zero copy overhead. + """ + return torch._C.ScriptDict(obj) # type: ignore[attr-defined] + + +def create_script_list(obj, type_hint=None): + """ + Create a ``torch._C.ScriptList`` instance with the data from ``obj``. + + Args: + obj (dict): The Python list that is used to initialize the ``ScriptList`` + returned by this function. + Returns: + An instance of ``torch._C.ScriptList`` that has the same data as ``obj`` + and can be passed between Python and TorchScript with reference semantics and + zero copy overhead. + """ + return torch._C.ScriptList(obj) # type: ignore[attr-defined] + + +_TOPLEVEL: bool = True + + +def _script_impl( + obj, + optimize=None, + _frames_up=0, + _rcb=None, + example_inputs: Union[List[Tuple], Dict[Callable, List[Tuple]], None] = None, +): + global type_trace_db + + if optimize is not None: + warnings.warn( + "`optimize` is deprecated and has no effect. " + "Use `with torch.jit.optimized_execution()` instead", + FutureWarning, + stacklevel=3, + ) + + # No-op for modules, functions, class instances that are already scripted + if isinstance(obj, RecursiveScriptClass): + return obj + if isinstance(obj, ScriptModule): + return obj + if isinstance(obj, ScriptFunction): + return obj + + if example_inputs: + # If MonkeyType is installed, enable profile directed type annotation + # Check if example_inputs are defined and generate call traces + # for the method by running eager mode version of the method with + # the provide example inputs. This logs all the traces in type_trace_db + type_trace_db = JitTypeTraceStore() + if monkeytype_trace: + monkeytype_config = JitTypeTraceConfig(type_trace_db) + with monkeytype_trace(monkeytype_config): + if isinstance(example_inputs, Dict): + # If the obj is an nn.Module or a class, then each method is + # executed with the arguments provided in the example inputs. + # example inputs here will be of type Dict(class.method, (arguments)) + # This is used to infer type annotations for those methods + # which are not called directly under the hood of monkeytype. + for module, example_input in example_inputs.items(): + for example in example_input: + module(*example) + elif isinstance(example_inputs, List): + for examples in example_inputs: + obj(*examples) + else: + raise ValueError( + "Error: Unable to infer types. Please format the inputs to type `List[Tuple]`" + " or `Dict[Callable, List[Tuple]]` to be run with MonkeyType." + ) + else: + warnings.warn( + "Warning: monkeytype is not installed. Please install https://github.com/Instagram/MonkeyType " + "to enable Profile-Directed Typing in TorchScript. Refer to " + "https://github.com/Instagram/MonkeyType/blob/master/README.rst to install MonkeyType. " + ) + + if isinstance(obj, torch.nn.Module): + obj = call_prepare_scriptable_func(obj) + return torch.jit._recursive.create_script_module( + obj, torch.jit._recursive.infer_methods_to_compile + ) + else: + obj = obj.__prepare_scriptable__() if hasattr(obj, "__prepare_scriptable__") else obj # type: ignore[operator] + + if isinstance(obj, dict): + return create_script_dict(obj) + if isinstance(obj, list): + return create_script_list(obj) + + if inspect.isclass(obj): + qualified_name = _qualified_name(obj) + # If this type is a `nn.Module` subclass, they probably meant to pass + # an instance instead of a Module + if issubclass(obj, torch.nn.Module): + raise RuntimeError( + f"Type '{obj}' cannot be compiled since it inherits from nn.Module, pass an instance instead" + ) + + # Enums are automatically usable in TorchScript, explicitly scripting + # is not necessary, but not harmful either. + if issubclass(obj, enum.Enum): + return obj + + if not _is_new_style_class(obj): + raise RuntimeError( + "TorchScript classes must be new-style classes. " + "Please inherit from 'object'." + ) + if len(obj.mro()) > 2: + raise RuntimeError( + "TorchScript classes does not support inheritance yet. " + "Please directly inherit from 'object'." + ) + if _rcb is None: + _rcb = _jit_internal.createResolutionCallbackFromFrame(_frames_up + 1) + _compile_and_register_class(obj, _rcb, qualified_name) + return obj + elif inspect.isfunction(obj) or inspect.ismethod(obj): + qualified_name = _qualified_name(obj) + # this is a decorated fn, and we need to the underlying fn and its rcb + if hasattr(obj, "__script_if_tracing_wrapper"): + obj = obj.__original_fn # type: ignore[union-attr] + _rcb = _jit_internal.createResolutionCallbackFromClosure(obj) + + # some functions are explicitly marked as not supported in script mode + if hasattr(obj, "__script_unsupported"): + raise RuntimeError("TorchScript error: " + obj.__script_unsupported) + + _check_directly_compile_overloaded(obj) + maybe_already_compiled_fn = _try_get_jit_cached_function(obj) + if maybe_already_compiled_fn: + maybe_already_compiled_fn._torchdynamo_inline = obj # type: ignore[attr-defined] + return maybe_already_compiled_fn + ast = get_jit_def(obj, obj.__name__) + if _rcb is None: + _rcb = _jit_internal.createResolutionCallbackFromClosure(obj) + fn = torch._C._jit_script_compile( + qualified_name, ast, _rcb, get_default_args(obj) + ) + # Forward docstrings + fn.__doc__ = obj.__doc__ + # Allow torch.compile() to inline + fn._torchdynamo_inline = obj # type: ignore[attr-defined] + _set_jit_function_cache(obj, fn) + return fn + else: + return torch.jit._recursive.create_script_class(obj) + + +def script( + obj, + optimize=None, + _frames_up=0, + _rcb=None, + example_inputs: Union[List[Tuple], Dict[Callable, List[Tuple]], None] = None, +): + r"""Script the function. + + Scripting a function or ``nn.Module`` will inspect the source code, compile + it as TorchScript code using the TorchScript compiler, and return a :class:`ScriptModule` or + :class:`ScriptFunction`. TorchScript itself is a subset of the Python language, so not all + features in Python work, but we provide enough functionality to compute on + tensors and do control-dependent operations. For a complete guide, see the + :ref:`language-reference`. + + Scripting a dictionary or list copies the data inside it into a TorchScript instance than can be + subsequently passed by reference between Python and TorchScript with zero copy overhead. + + ``torch.jit.script`` can be used as a function for modules, functions, dictionaries and lists + and as a decorator ``@torch.jit.script`` for :ref:`torchscript-classes` and functions. + + Args: + obj (Callable, class, or nn.Module): The ``nn.Module``, function, class type, + dictionary, or list to compile. + example_inputs (Union[List[Tuple], Dict[Callable, List[Tuple]], None]): Provide example inputs + to annotate the arguments for a function or ``nn.Module``. + + Returns: + If ``obj`` is ``nn.Module``, ``script`` returns + a :class:`ScriptModule` object. The returned :class:`ScriptModule` will + have the same set of sub-modules and parameters as the + original ``nn.Module``. If ``obj`` is a standalone function, + a :class:`ScriptFunction` will be returned. If ``obj`` is a ``dict``, then + ``script`` returns an instance of `torch._C.ScriptDict`. If ``obj`` is a ``list``, + then ``script`` returns an instance of `torch._C.ScriptList`. + + **Scripting a function** + The ``@torch.jit.script`` decorator will construct a :class:`ScriptFunction` + by compiling the body of the function. + + Example (scripting a function): + + .. testcode:: + + import torch + + @torch.jit.script + def foo(x, y): + if x.max() > y.max(): + r = x + else: + r = y + return r + + print(type(foo)) # torch.jit.ScriptFunction + + # See the compiled graph as Python code + print(foo.code) + + # Call the function using the TorchScript interpreter + foo(torch.ones(2, 2), torch.ones(2, 2)) + + .. testoutput:: + :hide: + + ... + + ****Scripting a function using example_inputs** + Example inputs can be used to annotate a function arguments. + + Example (annotating a function before scripting): + + .. testcode:: + + import torch + + def test_sum(a, b): + return a + b + + # Annotate the arguments to be int + scripted_fn = torch.jit.script(test_sum, example_inputs=[(3, 4)]) + + print(type(scripted_fn)) # torch.jit.ScriptFunction + + # See the compiled graph as Python code + print(scripted_fn.code) + + # Call the function using the TorchScript interpreter + scripted_fn(20, 100) + + .. testoutput:: + :hide: + + ... + + **Scripting an nn.Module** + Scripting an ``nn.Module`` by default will compile the ``forward`` method and recursively + compile any methods, submodules, and functions called by ``forward``. If a ``nn.Module`` only uses + features supported in TorchScript, no changes to the original module code should be necessary. ``script`` + will construct :class:`ScriptModule` that has copies of the attributes, parameters, and methods of + the original module. + + Example (scripting a simple module with a Parameter): + + .. testcode:: + + import torch + + class MyModule(torch.nn.Module): + def __init__(self, N, M): + super().__init__() + # This parameter will be copied to the new ScriptModule + self.weight = torch.nn.Parameter(torch.rand(N, M)) + + # When this submodule is used, it will be compiled + self.linear = torch.nn.Linear(N, M) + + def forward(self, input): + output = self.weight.mv(input) + + # This calls the `forward` method of the `nn.Linear` module, which will + # cause the `self.linear` submodule to be compiled to a `ScriptModule` here + output = self.linear(output) + return output + + scripted_module = torch.jit.script(MyModule(2, 3)) + + Example (scripting a module with traced submodules): + + .. testcode:: + + import torch + import torch.nn as nn + import torch.nn.functional as F + + class MyModule(nn.Module): + def __init__(self) -> None: + super().__init__() + # torch.jit.trace produces a ScriptModule's conv1 and conv2 + self.conv1 = torch.jit.trace(nn.Conv2d(1, 20, 5), torch.rand(1, 1, 16, 16)) + self.conv2 = torch.jit.trace(nn.Conv2d(20, 20, 5), torch.rand(1, 20, 16, 16)) + + def forward(self, input): + input = F.relu(self.conv1(input)) + input = F.relu(self.conv2(input)) + return input + + scripted_module = torch.jit.script(MyModule()) + + To compile a method other than ``forward`` (and recursively compile anything it calls), add + the :func:`@torch.jit.export ` decorator to the method. To opt out of compilation + use :func:`@torch.jit.ignore ` or :func:`@torch.jit.unused `. + + Example (an exported and ignored method in a module):: + + import torch + import torch.nn as nn + + class MyModule(nn.Module): + def __init__(self) -> None: + super().__init__() + + @torch.jit.export + def some_entry_point(self, input): + return input + 10 + + @torch.jit.ignore + def python_only_fn(self, input): + # This function won't be compiled, so any + # Python APIs can be used + import pdb + pdb.set_trace() + + def forward(self, input): + if self.training: + self.python_only_fn(input) + return input * 99 + + scripted_module = torch.jit.script(MyModule()) + print(scripted_module.some_entry_point(torch.randn(2, 2))) + print(scripted_module(torch.randn(2, 2))) + + Example ( Annotating forward of nn.Module using example_inputs):: + + import torch + import torch.nn as nn + from typing import NamedTuple + + class MyModule(NamedTuple): + result: List[int] + + class TestNNModule(torch.nn.Module): + def forward(self, a) -> MyModule: + result = MyModule(result=a) + return result + + pdt_model = TestNNModule() + + # Runs the pdt_model in eager model with the inputs provided and annotates the arguments of forward + scripted_model = torch.jit.script(pdt_model, example_inputs={pdt_model: [([10, 20, ], ), ], }) + + # Run the scripted_model with actual inputs + print(scripted_model([20])) + """ + if not _enabled: + return obj + try: + global _TOPLEVEL + prev = _TOPLEVEL + _TOPLEVEL = False + ret = _script_impl( + obj=obj, + optimize=optimize, + _frames_up=_frames_up + 1, + _rcb=_rcb, + example_inputs=example_inputs, + ) + + if prev: + log_torchscript_usage("script", model_id=_get_model_id(ret)) + + return ret + finally: + _TOPLEVEL = prev + + +# overloads are registered in _jit_internal and compiled here so that _overload +# can be used in nn/functional.py without an import cycle + + +def _check_overload_defaults(impl_defaults, overload_defaults, loc): + for name, overload_value in overload_defaults.items(): + if name not in impl_defaults or impl_defaults[name] != overload_value: + raise torch.jit.frontend.FrontendError( + loc, + "Default parameters on overloads do not affect the runtime so they " + "must equal to the default parameter on the implementation function. Found on " + f"parameter {name}", + ) + + +def _compile_function_with_overload(overload_fn, qual_name, impl_fn): + overload_decl = get_jit_def(overload_fn, overload_fn.__name__).decl() + overload_signature = torch.jit.annotations.get_signature( + overload_fn, None, None, inspect.ismethod(overload_fn) + ) + impl_ast = get_jit_def(impl_fn, impl_fn.__name__) + overload_defaults = get_default_args(overload_fn) + implementation_defaults = get_default_args(impl_fn) + _rcb = _jit_internal.createResolutionCallbackFromClosure(impl_fn) + _check_overload_defaults( + implementation_defaults, overload_defaults, overload_decl.range() + ) + fn = torch._C._jit_script_compile_overload( + qual_name, + overload_decl, + impl_ast, + _rcb, + implementation_defaults, + overload_signature, + ) + return fn + + +def _get_overloads(obj): + # check for cached compiled fns + existing_compiled_fns = _try_get_jit_cached_overloads(obj) + qual_name = _qualified_name(obj) + uncompiled_overloads = _jit_internal._get_fn_overloads(qual_name) + if uncompiled_overloads is None: + return existing_compiled_fns + + if obj in uncompiled_overloads: + raise RuntimeError( + _jit_internal.get_overload_no_implementation_error_message("function", obj) + ) + + compiled_fns = [ + _compile_function_with_overload(overload_fn, qual_name, obj) + for overload_fn in uncompiled_overloads + ] + + if existing_compiled_fns: + compiled_fns = existing_compiled_fns + compiled_fns + + # cache compilation, remove information stored to do compilation + _set_jit_overload_cache(obj, compiled_fns) + _jit_internal._clear_fn_overloads(qual_name) + return compiled_fns + + +def _check_directly_compile_overloaded(obj): + qual_name = _qualified_name(obj) + if _jit_internal._get_fn_overloads(qual_name) or _try_get_jit_cached_overloads(obj): + raise RuntimeError( + f"Function {qual_name} cannot be directly compiled because it" + " is overloaded. It must be used in a context of a function" + " where its inputs can determine which overload to call." + ) + + +def interface(obj): + r"""Decorate to annotate classes or modules of different types. + + This decorator can be used to define an interface that can be used to annotate + classes or modules of different types. This can be used for to annotate a submodule + or attribute class that could have different types that implement the same + interface, or which could be swapped at runtime; or to store a list of modules or + classes of varying types. + + It is sometimes used to implement "Callables" - functions or modules that implement + an interface but whose implementations differ and which can be swapped out. + + Example: + .. testcode:: + + import torch + from typing import List + + @torch.jit.interface + class InterfaceType: + def run(self, x: torch.Tensor) -> torch.Tensor: + pass + + # implements InterfaceType + @torch.jit.script + class Impl1: + def run(self, x: torch.Tensor) -> torch.Tensor: + return x.relu() + + class Impl2(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.val = torch.rand(()) + + @torch.jit.export + def run(self, x: torch.Tensor) -> torch.Tensor: + return x + self.val + + def user_fn(impls: List[InterfaceType], idx: int, val: torch.Tensor) -> torch.Tensor: + return impls[idx].run(val) + + user_fn_jit = torch.jit.script(user_fn) + + impls = [Impl1(), torch.jit.script(Impl2())] + val = torch.rand(4, 4) + user_fn_jit(impls, 0, val) + user_fn_jit(impls, 1, val) + """ + if not inspect.isclass(obj): + raise RuntimeError("interface must be applied to a class") + if not _is_new_style_class(obj): + raise RuntimeError("TorchScript interfaces must inherit from 'object'") + + # Expected MRO is: + # User module + # torch.nn.modules.module.Module + # object + is_module_interface = issubclass(obj, torch.nn.Module) and len(obj.mro()) == 3 + + if not is_module_interface and len(obj.mro()) > 2: + raise RuntimeError( + "TorchScript interface does not support inheritance yet. " + "Please directly inherit from 'object' or 'nn.Module'." + ) + + qualified_name = _qualified_name(obj) + rcb = _jit_internal.createResolutionCallbackFromFrame(1) + # if this type is a `nn.Module` subclass, generate a module interface type + # instead of a class interface type; a module interface type only compiles + # the user provided methods as part of the interface + ast = get_jit_class_def(obj, obj.__name__) + mangled_classname = torch._C._jit_script_interface_compile( + qualified_name, ast, rcb, is_module_interface + ) + obj.__torch_script_interface__ = mangled_classname + return obj + + +def _recursive_compile_class(obj, loc): + _qual_name = _qualified_name(obj) + # We're starting a new compilation, so update the error call stack in + # case it fails + error_stack = torch._C.CallStack(_qual_name, loc) # noqa: F841 + rcb = _jit_internal.createResolutionCallbackForClassMethods(obj) + return _compile_and_register_class(obj, rcb, _qual_name) + + +CompilationUnit = torch._C.CompilationUnit +set_module(CompilationUnit, "torch.jit") + + +def pad(s: str, padding: int, offset: int = 0, char: str = " "): + if padding >= len(s): + padding -= len(s) + return "".join([char for _ in range(padding + offset)]) + s + + +class _ScriptProfileColumn: + def __init__(self, header: str, alignment: int = 4, offset: int = 0): + self.header = header + self.alignment = alignment + self.offset = offset + self.rows: Dict[int, Any] = {} + + def add_row(self, lineno: int, value: Any): + self.rows[lineno] = value + + def materialize(self): + max_length = len(self.header) + rows: List[Tuple[int, str]] = [] + for key, value in self.rows.items(): + cell = str(value) + rows.append((key, cell)) + max_length = max(len(cell), max_length) + + if self.alignment > 0: + padding = max_length + self.alignment + padding -= padding % self.alignment + else: + padding = 0 + + rows = [(key, pad(cell, padding, self.offset)) for key, cell in rows] + return pad(self.header, padding, self.offset), rows + + +class _ScriptProfileTable: + def __init__(self, cols: List[_ScriptProfileColumn], source_range: List[int]): + self.cols = cols + self.source_range = source_range + + def dump_string(self): + outputs: List[str] = [] + cells: List[Tuple[str, Dict[int, str]]] = [] + header_buffer = "" + for col in self.cols: + header, rows = col.materialize() + header_buffer += header + cells.append((header, dict(rows))) + + outputs.append(header_buffer) + outputs.append(pad("", len(header_buffer), 0, "=")) + for line in self.source_range: + row_buffer = "" + for header, rows in cells: + cell = rows.get(line) + if cell is None: + row_buffer += pad("", len(header)) + else: + row_buffer += cell + outputs.append(row_buffer) + return "\n".join(outputs) + + +class _ScriptProfile: + def __init__(self) -> None: + self.profile = classes.profiling._ScriptProfile() + + def enable(self): + self.profile.enable() + + def disable(self): + self.profile.disable() + + def dump_string(self) -> str: + outputs: List[str] = [] + for source_stats in self.profile._dump_stats(): + source_ref = source_stats.source() + source_lines = source_ref.text().splitlines() + dedent = min(len(line) - len(line.lstrip(" ")) for line in source_lines) + source_lines = [line[dedent:] for line in source_lines] + + start_line = source_ref.starting_lineno() + end_line = start_line + len(source_lines) + source_range = range(start_line, end_line) + lineno = _ScriptProfileColumn("Line #") + hits = _ScriptProfileColumn("Hits") + time_ns = _ScriptProfileColumn("Time (ns)") + line_contents = _ScriptProfileColumn("Line Contents", 0, 1) + stats = source_stats.line_map() + for line in source_range: + lineno.add_row(line, line) + line_contents.add_row(line, source_lines[line - start_line]) + stat = stats.get(line) + if stat is not None: + hits.add_row(line, stat.count()) + time_ns.add_row(line, stat.duration_ns()) + + table = _ScriptProfileTable( + [lineno, hits, time_ns, line_contents], list(source_range) + ) + outputs.append(table.dump_string()) + return "\n\n".join(outputs) + + def dump(self): + print(self.dump_string()) + + +def _unwrap_optional(x): + assert x is not None, "Unwrapping null optional" + return x + + +_register_builtin(_unwrap_optional, "aten::_unwrap_optional") +_register_builtin(_jit_internal.is_scripting, "aten::is_scripting") +_register_builtin(has_torch_function, "aten::has_torch_function") +_register_builtin(has_torch_function_unary, "aten::has_torch_function") +_register_builtin(has_torch_function_variadic, "aten::has_torch_function") diff --git a/lib/python3.12/site-packages/torch/jit/_serialization.py b/lib/python3.12/site-packages/torch/jit/_serialization.py new file mode 100644 index 0000000000000000000000000000000000000000..a2cf33e544701df489f056526bdd91b1b128c826 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_serialization.py @@ -0,0 +1,273 @@ +# mypy: allow-untyped-defs +"""Serialization. + +This module contains functionality for serializing TorchScript modules, notably: + * torch.jit.save + * torch.jit.load + +This is not intended to be imported directly; please use the exposed +functionalities in `torch.jit`. +""" + +import os + +import torch +from torch._jit_internal import _get_model_id +from torch._utils_internal import log_torchscript_usage +from torch.jit._recursive import wrap_cpp_module +from torch.serialization import validate_cuda_device + + +def save(m, f, _extra_files=None): + r""" + Save an offline version of this module for use in a separate process. + + The saved module serializes all of the methods, submodules, parameters, and + attributes of this module. It can be loaded into the C++ API using + ``torch::jit::load(filename)`` or into the Python API with + :func:`torch.jit.load `. + + To be able to save a module, it must not make any calls to native Python + functions. This means that all submodules must be subclasses of + :class:`ScriptModule` as well. + + .. DANGER:: + All modules, no matter their device, are always loaded onto the CPU + during loading. This is different from :func:`torch.load`'s semantics + and may change in the future. + + Args: + m: A :class:`ScriptModule` to save. + f: A file-like object (has to implement write and flush) or a string + containing a file name. + _extra_files: Map from filename to contents which will be stored as part of `f`. + + .. note:: + torch.jit.save attempts to preserve the behavior of some operators + across versions. For example, dividing two integer tensors in + PyTorch 1.5 performed floor division, and if the module + containing that code is saved in PyTorch 1.5 and loaded in PyTorch 1.6 + its division behavior will be preserved. The same module saved in + PyTorch 1.6 will fail to load in PyTorch 1.5, however, since the + behavior of division changed in 1.6, and 1.5 does not know how to + replicate the 1.6 behavior. + + Example: + .. testcode:: + + import torch + import io + + class MyModule(torch.nn.Module): + def forward(self, x): + return x + 10 + + m = torch.jit.script(MyModule()) + + # Save to file + torch.jit.save(m, 'scriptmodule.pt') + # This line is equivalent to the previous + m.save("scriptmodule.pt") + + # Save to io.BytesIO buffer + buffer = io.BytesIO() + torch.jit.save(m, buffer) + + # Save with extra files + extra_files = {'foo.txt': b'bar'} + torch.jit.save(m, 'scriptmodule.pt', _extra_files=extra_files) + """ + log_torchscript_usage("save", model_id=_get_model_id(m)) + if _extra_files is None: + _extra_files = {} + if isinstance(f, (str, os.PathLike)): + m.save(f, _extra_files=_extra_files) + else: + ret = m.save_to_buffer(_extra_files=_extra_files) + f.write(ret) + + +def load(f, map_location=None, _extra_files=None, _restore_shapes=False): + r""" + Load a :class:`ScriptModule` or :class:`ScriptFunction` previously saved with :func:`torch.jit.save `. + + All previously saved modules, no matter their device, are first loaded onto CPU, + and then are moved to the devices they were saved from. If this fails (e.g. + because the run time system doesn't have certain devices), an exception is + raised. + + Args: + f: a file-like object (has to implement read, readline, tell, and seek), + or a string containing a file name + map_location (string or torch.device): A simplified version of + ``map_location`` in `torch.jit.save` used to dynamically remap + storages to an alternative set of devices. + _extra_files (dictionary of filename to content): The extra + filenames given in the map would be loaded and their content + would be stored in the provided map. + _restore_shapes (bool): Whether or not to retrace the module on load using stored inputs + + Returns: + A :class:`ScriptModule` object. + + Example: + .. testcode:: + + import torch + import io + + torch.jit.load('scriptmodule.pt') + + # Load ScriptModule from io.BytesIO object + with open('scriptmodule.pt', 'rb') as f: + buffer = io.BytesIO(f.read()) + + # Load all tensors to the original device + torch.jit.load(buffer) + + # Load all tensors onto CPU, using a device + buffer.seek(0) + torch.jit.load(buffer, map_location=torch.device('cpu')) + + # Load all tensors onto CPU, using a string + buffer.seek(0) + torch.jit.load(buffer, map_location='cpu') + + # Load with extra files. + extra_files = {'foo.txt': ''} # values will be replaced with data + torch.jit.load('scriptmodule.pt', _extra_files=extra_files) + print(extra_files['foo.txt']) + + .. testoutput:: + :hide: + + ... + + .. testcleanup:: + + import os + os.remove("scriptmodule.pt") + """ + if isinstance(f, (str, os.PathLike)): + if not os.path.exists(f): + raise ValueError(f"The provided filename {f} does not exist") + if os.path.isdir(f): + raise ValueError(f"The provided filename {f} is a directory") + + map_location = validate_map_location(map_location) + if _extra_files is None: + _extra_files = {} + + cu = torch._C.CompilationUnit() + if isinstance(f, (str, os.PathLike)): + cpp_module = torch._C.import_ir_module(cu, os.fspath(f), map_location, _extra_files, _restore_shapes) # type: ignore[call-arg] + else: + cpp_module = torch._C.import_ir_module_from_buffer( + cu, f.read(), map_location, _extra_files, _restore_shapes + ) # type: ignore[call-arg] + + # TODO: Pretty sure this approach loses ConstSequential status and such + ret = wrap_cpp_module(cpp_module) + log_torchscript_usage("load", model_id=_get_model_id(ret)) + return ret + + +def validate_map_location(map_location=None): + if isinstance(map_location, str): + map_location = torch.device(map_location) + elif not (map_location is None or isinstance(map_location, torch.device)): + raise ValueError( + "map_location should be either None, string or torch.device, " + "but got type: " + str(type(map_location)) + ) + + if str(map_location).startswith("cuda"): + validate_cuda_device(map_location) + + return map_location + + +def jit_module_from_flatbuffer(f): + if isinstance(f, (str, os.PathLike)): + f = os.fspath(f) + return wrap_cpp_module(torch._C._load_jit_module_from_file(f)) + else: + return wrap_cpp_module(torch._C._load_jit_module_from_bytes(f.read())) + + +def save_jit_module_to_flatbuffer(m, f, _extra_files=None): + r""" + Save an offline version of this module for use in a separate process. + + The saved module serializes all of the methods, submodules, parameters, and + attributes of this module. It can be loaded into the C++ API using + ``torch::jit::load_jit_module_from_file(filename)`` or into the Python API with + :func:`torch.jit.jit_module_from_flatbuffer`. + + To be able to save a module, it must not make any calls to native Python + functions. This means that all submodules must be subclasses of + :class:`ScriptModule` as well. + + .. DANGER:: + All modules, no matter their device, are always loaded onto the CPU + during loading. This is different from :func:`torch.load`'s semantics + and may change in the future. + + Args: + m: A :class:`ScriptModule` to save. + f: A string for file path + + + Example: + .. testcode:: + + import torch + import io + + class MyModule(torch.nn.Module): + def forward(self, x): + return x + 10 + + m = torch.jit.script(MyModule()) + + # Save to file + torch.jit.save_jit_module_to_flatbuffer(m, 'scriptmodule.ff') + """ + extra_files = _extra_files + if extra_files is None: + extra_files = {} + + if isinstance(f, (str, os.PathLike)): + f = os.fspath(f) + torch._C._save_jit_module(m._c, f, extra_files) + else: + s = torch._C._save_jit_module_to_bytes(m._c, extra_files) + f.write(s) + + +def get_flatbuffer_module_info(path_or_file): + r"""Get some information regarding a model file in flatbuffer format. + + Args: + path_or_file: Either str, Path or file like object (BytesIO OK). + If it's str or Path, we will read the file referenced by that + path as Bytes. + + Returns: + A dict with metadata on what that file contains, currently looks like + this: + { + 'bytecode_version': 4, # int + 'operator_version': 4, # int + 'function_names': { + '__torch__.___torch_mangle_0.Foo.forward'}, # set + 'type_names': set(), # set + 'opname_to_num_args': {'aten::linear': 3} # Dict[str, int] + } + """ + if isinstance(path_or_file, (str, os.PathLike)): + with open(path_or_file, "rb") as f: + all_bytes = f.read() + else: + all_bytes = path_or_file.read() + return torch._C._get_module_info_from_flatbuffer(all_bytes) diff --git a/lib/python3.12/site-packages/torch/jit/_trace.py b/lib/python3.12/site-packages/torch/jit/_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..ef5292fe93ecb8b5c55dde96579bb7a4186d2ab8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/_trace.py @@ -0,0 +1,1501 @@ +# mypy: allow-untyped-defs +"""Tracing. + +This module contains functionality to support the JIT's tracing frontend, notably: + * torch.jit.trace + * torch.jit.trace_module + +This is not intended to be imported directly; please use the exposed +functionalities in `torch.jit`. +""" + +import contextlib +import copy +import functools +import inspect +import os +import re +import warnings +from enum import Enum +from typing import Any, Callable, Dict, List, Optional, Set, TypeVar +from typing_extensions import ParamSpec + +import torch +from torch._jit_internal import ( + _get_model_id, + _qualified_name, + get_callable_argument_names, + is_scripting, +) +from torch.autograd import function +from torch.jit._script import _CachedForward, script, ScriptModule +from torch.jit._state import _enabled, _python_cu +from torch.nn import Module +from torch.testing._comparison import default_tolerances + + +_flatten = torch._C._jit_flatten +_unflatten = torch._C._jit_unflatten + +R = TypeVar("R", covariant=True) # return type (always covariant) +P = ParamSpec("P") + + +def _create_interpreter_name_lookup_fn(frames_up=1): + def _get_interpreter_name_for_var(var): + frame = inspect.currentframe() + if not frame: + raise RuntimeError("failed to inspect frame") + + i = 0 + while i < frames_up + 1: + frame = frame.f_back + if not frame: + raise RuntimeError("failed to get frame") + i += 1 + + f_locals = frame.f_locals + + for k, v in f_locals.items(): + if isinstance(v, torch.Tensor) and var is v: + return k if k != "self" else "" + return "" + + return _get_interpreter_name_for_var + + +def _unique_state_dict(module, keep_vars=False): + # since Parameter.detach() always creates a new torch.Tensor instance, + # id(v) doesn't work with it. So we always get the Parameter or Buffer + # as values, and deduplicate the params using Parameters and Buffers + state_dict = module.state_dict(keep_vars=True) + filtered_dict = type(state_dict)() + seen_ids: Set[int] = set() + for k, v in state_dict.items(): + if id(v) in seen_ids: + continue + seen_ids.add(id(v)) + if keep_vars: + filtered_dict[k] = v + else: + filtered_dict[k] = v.detach() + return filtered_dict + + +class ONNXTracedModule(torch.nn.Module): + def __init__( + self, + inner, + strict=True, + force_outplace=False, + return_inputs=False, + return_inputs_states=False, + ): + super().__init__() + # inner may be a Module, or it may be an arbitrary callable + # If it's a Module, we get its parameters automatically, which lets + # us avoid a special casing functions versus modules. + self.inner = inner + self.strict = strict + self._force_outplace = force_outplace + self._return_inputs = return_inputs + self._return_inputs_states = return_inputs_states + + def forward(self, *args: torch.Tensor): + in_vars, in_desc = _flatten(args) + # NOTE: use full state, because we need it for BatchNorm export + # This differs from the compiler path, which doesn't support it at the moment. + module_state = list(_unique_state_dict(self, keep_vars=True).values()) + + ret_inputs = [] + inputs_states = [] + outs = [] + + def wrapper(*args): + in_args: List[torch.Tensor] = [] + for i in range(len(in_vars)): + if not isinstance(args[i], torch.Tensor): + raise RuntimeError("Expected Tensor argument") + in_args.append(args[i]) + + trace_inputs = _unflatten(in_args, in_desc) + + if self._return_inputs: + ret_inputs.append( + tuple(x.clone(memory_format=torch.preserve_format) for x in args) + ) + if self._return_inputs_states: + inputs_states.append(_unflatten(in_args, in_desc)) + outs.append(self.inner(*trace_inputs)) + if self._return_inputs_states: + inputs_states[0] = (inputs_states[0], trace_inputs) + out_vars, _ = _flatten(outs) + if len(out_vars) == 1: + return out_vars[0] + else: + return tuple(out_vars) + + graph, _out = torch._C._create_graph_by_tracing( + wrapper, + in_vars + module_state, + _create_interpreter_name_lookup_fn(), + self.strict, + self._force_outplace, + ) + + if self._return_inputs: + return graph, outs[0], ret_inputs[0] + if self._return_inputs_states: + return graph, outs[0], inputs_states[0] + else: + return graph, outs[0] + + +def _clone_inputs(args): + def clone_input(a): + if a is None: + return None + elif isinstance(a, torch.Tensor): + # TODO: figure out one liner to .clone() and set requires_grad + v = ( + a.detach() + .clone(memory_format=None if a.is_mkldnn else torch.preserve_format) + .requires_grad_(a.requires_grad) + ) + if a.grad is not None: + v.grad = clone_input(v.grad) + return v + else: + return a.clone(memory_format=torch.preserve_format) + + return function._nested_map( + lambda x: isinstance(x, torch.Tensor), clone_input, condition_msg="tensors" + )(args) + + +# This is purely for developer debugging. We are not going to advertise it. +_JIT_TIME = os.environ.get("PYTORCH_JIT_TIME", False) # CUDA-only timing +_JIT_DISABLE = os.environ.get("PYTORCH_JIT_DISABLE", False) +_JIT_STATS = os.environ.get("PYTORCH_JIT_STATS", False) + + +@contextlib.contextmanager +def _time(trace_name, name, time=True): + if (not _JIT_TIME and not time) or not torch.cuda.is_available(): + yield + return + stream = torch.cuda.current_stream() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + stream.record_event(start) + try: + yield + finally: + stream.record_event(end) + end.synchronize() + print(f"{trace_name} {name} time: {start.elapsed_time(end)} ms") + + +def verify(model, args, loss_fn=torch.sum, devices=None): + """ + Verify that a JIT compiled model has the same behavior as its uncompiled version along with its backwards pass. + + If your model returns multiple outputs, + you must also specify a `loss_fn` to produce a loss for which + the backwards will be computed. + + This function has side-effects (e.g., it executes your model / saves and loads + parameters), so don't expect the model to come out exactly the same as what + you passed in. + + Args: + model (compiled torch.nn.Module or function): the module/function to be + verified. The module/function definition MUST have been decorated with + `@torch.jit.compile`. + args (tuple or Tensor): the positional arguments to pass to the + compiled function/module to be verified. A non-tuple is assumed to + be a single positional argument to be passed to the model. + loss_fn (function, optional): the loss function to be applied to + the output of the model, before backwards is invoked. By default, + we assume that a model returns a single result, and we :func:`torch.sum` + before calling backwards; if this is inappropriate, you can pass your + own loss function. Note that if a model returns a tuple of results, + these are passed as separate positional arguments to `loss_fn`. + devices (iterable of device IDs, optional): the GPU devices which the + compiled module will be run on. This determines the RNG state we + must save when running both compiled and uncompiled versions of the model. + """ + # TODO: In principle, we track device information in our trace, so it + # should be possible to check if our execution actually obeyed the 'devices' + # the user provided. + + # TODO: Consider adding a utility function to torch.jit to test + # for this case + if not isinstance(model, torch._C.CompiledFunction): # type: ignore[attr-defined] + raise TypeError( + "Cannot verify an uncompiled module. Add @torch.jit.compile to compile it" + ) + is_module = isinstance(model, Module) + + if not isinstance(args, tuple): + args = (args,) + + if is_module: + saved_state = copy.deepcopy(model.state_dict()) + + def run_fwd_bwd(args, force_trace=False, assert_compiled=False): + params = list(model.parameters()) if is_module else [] + in_vars, _ = _flatten((args, params)) + # We use a special API to reset the trace and compile it from scratch. + compiled_fn = model + if force_trace: + compiled_fn.clear_cache() + if assert_compiled: + hits = compiled_fn.hits + out = model(*args) + if assert_compiled and compiled_fn.hits == hits: # type: ignore[possibly-undefined] + raise RuntimeError("failed to use the compiled function") + if not isinstance(out, tuple): + out = (out,) + if loss_fn == torch.sum and len(out) != 1: + raise ValueError( + f"Model returns {len(out)} outputs, but default loss function " + "(torch.sum) can only handle a single output" + ) + out_vars, _ = _flatten(out) + saved_outs = [ + v.detach().clone(memory_format=torch.preserve_format) for v in out_vars + ] + loss = loss_fn(*out) + grads = torch.autograd.grad([loss], in_vars) + # TODO: I'm not sure if the clone here is necessary but it is safer + saved_grads = [ + v.detach().clone(memory_format=torch.preserve_format) for v in grads + ] + return (saved_outs, saved_grads) + + with torch.random.fork_rng(devices, _caller="torch.jit.verify"): + uncompiled_outs, uncompiled_grads = run_fwd_bwd(args, force_trace=True) + assert model.has_trace_for(*args) + + if is_module: + model.load_state_dict(saved_state) # type: ignore[possibly-undefined] + compiled_outs, compiled_grads = run_fwd_bwd(args, assert_compiled=True) + + _verify_equal(uncompiled_outs, compiled_outs) + _verify_equal(uncompiled_grads, compiled_grads) + + +def _verify_equal(xs, ys): + for x, y in zip(xs, ys): + if x.sub(y).abs().max() > 1e-6: + raise RuntimeError("JIT and real computation mismatch") + + +def indent(s): + return "\n".join(["\t" + line for line in s.splitlines()]) + + +class TracingCheckError(Exception): + def __init__(self, graph_diff_error, tensor_compare_error, extra_msg=None): + self.message = "Tracing failed sanity checks!\n" + if extra_msg is not None: + self.message += extra_msg + "\n" + if graph_diff_error is not None: + self.message += "ERROR: Graphs differed across invocations!\n" + self.message += indent(graph_diff_error) + "\n" + if tensor_compare_error is not None: + self.message += ( + "ERROR: Tensor-valued Constant nodes differed in value " + "across invocations. This often indicates that the tracer has" + " encountered untraceable code.\n" + ) + self.message += indent(tensor_compare_error) + "\n" + super().__init__(self.message) + + +# Check the traced module against a set of user-provided validation inputs +@torch.no_grad() +def _check_trace( + check_inputs, + func, + traced_func, + check_tolerance, + strict, + force_outplace, + is_trace_module, + _module_class, + example_inputs_is_kwarg=False, +): + # Note: tracing is independent of optimizations, which consume the trace + for inputs in check_inputs: + if isinstance(inputs, torch.Tensor): + inputs = (inputs,) + + if is_trace_module: + copied_dict = {} + for name, data in inputs.items(): + copied_dict[name] = _clone_inputs(data) + check_mod = torch.jit.trace_module( + getattr(func, "__self__", func), + copied_dict, + check_trace=False, + strict=strict, + _force_outplace=force_outplace, + _module_class=_module_class, + _compilation_unit=torch._C.CompilationUnit(), + example_inputs_is_kwarg=example_inputs_is_kwarg, + _store_inputs=False, + ) + check_mod_func = check_mod._c._get_method(traced_func.name) + inputs = inputs[traced_func.name] + if ( + isinstance(inputs, (torch.Tensor)) + or isinstance(inputs, dict) + and not example_inputs_is_kwarg + ): + inputs = (inputs,) + else: + if example_inputs_is_kwarg: + check_mod = torch.jit.trace( + func, + check_trace=False, + strict=strict, + _force_outplace=force_outplace, + _module_class=_module_class, + example_kwarg_inputs=_clone_inputs(inputs), + _store_inputs=False, + ) + else: + check_mod = torch.jit.trace( + func, + _clone_inputs(inputs), + check_trace=False, + strict=strict, + _force_outplace=force_outplace, + _module_class=_module_class, + _store_inputs=False, + ) + check_mod_func = check_mod + + def graph_diagnostic_info(): + mod_canonicalized = torch._C._jit_pass_canonicalize(traced_func.graph) + torch._C._jit_pass_inline(mod_canonicalized) + torch._C._jit_pass_erase_shape_information(mod_canonicalized) + mod_str = str(mod_canonicalized) + mod_str = re.sub(r"___torch_mangle_[0-9]+\.", "", mod_str) + check_canonicalized = torch._C._jit_pass_canonicalize(check_mod_func.graph) + torch._C._jit_pass_inline(check_canonicalized) + torch._C._jit_pass_erase_shape_information(check_canonicalized) + check_str = str(check_canonicalized) + check_str = re.sub(r"___torch_mangle_[0-9]+\.", "", check_str) + + graph_diff_errors = None + if mod_str != check_str: + import difflib + + graph_diff = difflib.ndiff( + mod_str.splitlines(True), check_str.splitlines(True) + ) + graph_diff_errors = "Graph diff:\n" + indent("".join(graph_diff)) + "\n" + + for n_mod, n_check in zip( + mod_canonicalized.nodes(), check_canonicalized.nodes() + ): + if str(n_mod) != str(n_check): + graph_diff_errors += "First diverging operator:\n" + node_diff = difflib.ndiff( + str(n_mod).splitlines(True), str(n_check).splitlines(True) + ) + source_printout = ( + "Node diff:\n" + indent("".join(node_diff)) + "\n" + ) + mod_stack = n_mod.sourceRange() + if mod_stack: + source_printout += ( + "Trace source location:\n" + indent(mod_stack) + "\n" + ) + check_stack = n_check.sourceRange() + if check_stack: + source_printout += ( + "Check source location:\n" + indent(check_stack) + "\n" + ) + graph_diff_errors += source_printout + + break # For now, only print out the first pair of nodes that diverges + + tensor_compare_errors = None + # Check Tensor-valued constant nodes + for n_mod, n_check in zip( + mod_canonicalized.nodes(), check_canonicalized.nodes() + ): + if n_mod.kind() != n_check.kind(): + break # Graphs have already diverged + + if n_mod.kind() == "prim::Constant" and not ( + n_mod.mustBeNone() or n_check.mustBeNone() + ): + if not n_mod.hasAttribute("value"): + continue + if n_mod.kindOf("value") != "t" or n_check.kindOf("value") != "t": + continue + + mod_tensor_val = n_mod.t("value") + check_tensor_val = n_check.t("value") + + try: + torch.testing.assert_close( + mod_tensor_val, check_tensor_val, equal_nan=True + ) + except (RuntimeError, AssertionError) as e: + if tensor_compare_errors is None: + tensor_compare_errors = "" + tensor_compare_errors += "Node:\n" + indent(str(n_mod)) + "\n" + compare_stack = n_mod.sourceRange() + if compare_stack: + tensor_compare_errors += ( + "Source Location:\n" + indent(compare_stack) + "\n" + ) + tensor_compare_errors += "Comparison exception: " + indent( + str(e) + ) + + break # For now, only print the first diverging pair + + return graph_diff_errors, tensor_compare_errors + + def wrap_retval(x): + return x if isinstance(x, tuple) else (x,) + + def run_mod_and_filter_tensor_outputs(mod, inputs, running_what): + try: + if isinstance(inputs, dict) and example_inputs_is_kwarg: + outs = wrap_retval(mod(**inputs)) + else: + outs = wrap_retval(mod(*_clone_inputs(inputs))) + outs = [out for out in outs if isinstance(out, torch.Tensor)] + return outs + except Exception as e: + graph_diff_errors, tensor_compare_errors = graph_diagnostic_info() + msg = f"encountered an exception while running the {running_what} with test inputs.\nException:\n{indent(str(e))}" + raise TracingCheckError( + graph_diff_errors, + tensor_compare_errors, + extra_msg=msg, + ) from e + + has_warned = [False] + + def maybe_warn_nondeterministic(): + if has_warned[0]: + return + has_warned[0] = True + nondeterm_ops = [ + op for op in traced_func.graph.nodes() if op.isNondeterministic() + ] + if len(nondeterm_ops) > 0: + nondeterministic_ops_warning = "Trace had nondeterministic nodes. " + nondeterministic_ops_warning += ( + "Did you forget call .eval() on your model? Nodes:\n" + ) + nondeterministic_ops_warning += "\n".join( + [indent(str(op)) for op in nondeterm_ops][:20] + ) + nondeterministic_ops_warning += ( + "\nThis may cause errors in trace checking. To disable trace checking," + " pass check_trace=False to torch.jit.trace()" + ) + warnings.warn( + nondeterministic_ops_warning, category=TracerWarning, stacklevel=5 + ) + + def compare_outputs(original, reference, match_what): + all_ok = True + for i, (orig, ref) in enumerate(zip(original, reference)): + try: + if orig.is_quantized: + orig = orig.dequantize() + if ref.is_quantized: + ref = ref.dequantize() + if orig.is_mkldnn: + orig = orig.to_dense() + if ref.is_mkldnn: + ref = ref.to_dense() + if ref.is_complex() or orig.is_complex(): + torch.testing.assert_close( + orig.to(torch.cdouble), + ref.to(torch.cdouble), + rtol=check_tolerance, + atol=default_tolerances(orig, ref)[1], + equal_nan=True, + ) + else: + if orig.is_mps or ref.is_mps: + torch.testing.assert_close( + orig.float(), + ref.float(), + rtol=check_tolerance, + atol=default_tolerances(orig, ref)[1], + equal_nan=True, + ) + elif getattr(orig, "is_nested", None) or getattr( + ref, "is_nested", None + ): + assert getattr(orig, "is_nested", None) == getattr( + ref, "is_nested", None + ) + for t_orig, t_ref in zip(orig.unbind(), ref.unbind()): + torch.testing.assert_close( + t_orig.double(), + t_ref.double(), + rtol=check_tolerance, + atol=default_tolerances(t_orig, t_ref)[1], + equal_nan=True, + ) + else: + torch.testing.assert_close( + orig.double(), + ref.double(), + rtol=check_tolerance, + atol=default_tolerances(orig, ref)[1], + equal_nan=True, + ) + except AssertionError as e: + maybe_warn_nondeterministic() + warnings.warn( + "Output nr " + + str(i + 1) + + ". of the traced function does not match " + "the corresponding output of the " + + match_what + + ". Detailed error:\n" + + str(e), + category=TracerWarning, + stacklevel=4, + ) + all_ok = False + + return all_ok + + traced_outs = run_mod_and_filter_tensor_outputs(traced_func, inputs, "trace") + fn_outs = run_mod_and_filter_tensor_outputs(func, inputs, "Python function") + if compare_outputs(traced_outs, fn_outs, "Python function"): + check_outs = run_mod_and_filter_tensor_outputs( + check_mod_func, inputs, "repeated trace" + ) + compare_outputs(traced_outs, check_outs, "repeated trace") + + diag_info = graph_diagnostic_info() + if any(info is not None for info in diag_info): + raise TracingCheckError(*diag_info) + + +class TracerWarning(Warning): + @staticmethod + def ignore_lib_warnings(): + # We ignore warnings from all submodules excluding the JIT, because we need them e.g. for _check_trace + warnings.filterwarnings( + "ignore", category=TracerWarning, module="torch.(?!jit)" + ) + warnings.filterwarnings("ignore", "torch::jit::fuser::cuda") + + +# We ignore the tracer warnings coming form inside the library, because all our shape +# checks in nn will trigger them. +TracerWarning.ignore_lib_warnings() +torch._C._tracer_warn_use_python() + + +def make_tuple(example_inputs): + if isinstance(example_inputs, (torch.Tensor, dict)): + return (example_inputs,) + # done primarily so that weird iterables fail here and not pybind11 code + if not isinstance(example_inputs, tuple): + return tuple(example_inputs) + return example_inputs + + +def make_module(mod, _module_class, _compilation_unit): + if isinstance(mod, ScriptModule): + return mod + elif torch._jit_internal.module_has_exports(mod): + infer_methods_stubs_fn = torch.jit._recursive.make_stubs_from_exported_methods + return torch.jit._recursive.create_script_module( + mod, infer_methods_stubs_fn, share_types=False, is_tracing=True + ) + else: + if _module_class is None: + _module_class = TopLevelTracedModule + return _module_class(mod, _compilation_unit=_compilation_unit) + + +def wrap_check_inputs(check_inputs): + if check_inputs is None: + return None + + return [{"forward": c} for c in check_inputs] + + +def analyze_ts_result_with_export_result(export, trace): + import torch.utils._pytree as pytree + + flat_export = pytree.tree_leaves(export) + flat_trace = pytree.tree_leaves(trace) + + for orig, loaded in zip(flat_export, flat_trace): + if orig.layout != loaded.layout: + return False + # mkldnn is not supported for torch.allclose + if orig.layout == torch._mkldnn: # type: ignore[attr-defined] + return True + if type(orig) != type(loaded): + return False + + if isinstance(orig, torch._subclasses.FakeTensor): + # Skip for FakeTensor. + return True + elif isinstance(orig, torch.Tensor): + if orig.dtype != loaded.dtype: + return False + if not torch.allclose(orig, loaded): + return False + else: + if orig != loaded: + return False + return True + + +def _trace_impl( + func, + example_inputs=None, + optimize=None, + check_trace=True, + check_inputs=None, + check_tolerance=1e-5, + strict=True, + _force_outplace=False, + _module_class=None, + _compilation_unit=_python_cu, + example_kwarg_inputs=None, + _store_inputs=True, +): + if isinstance(func, torch.jit.ScriptModule): + # it is hard to trace it because the forward method on ScriptModule is already defined, so it + # would result in an error. + warnings.warn( + "The input to trace is already a ScriptModule, tracing it is a no-op. Returning the object as is." + ) + return func + + if isinstance(func, torch.nn.Module): + if example_inputs is None: + if isinstance(example_kwarg_inputs, dict): + example_inputs = example_kwarg_inputs + else: + raise RuntimeError("example_kwarg_inputs should be a dict") + return trace_module( + func, + {"forward": example_inputs}, + None, + check_trace, + wrap_check_inputs(check_inputs), + check_tolerance, + strict, + _force_outplace, + _module_class, + example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict), + _store_inputs=_store_inputs, + ) + if ( + hasattr(func, "__self__") + and isinstance(func.__self__, torch.nn.Module) + and func.__name__ == "forward" + ): + if example_inputs is None: + if isinstance(example_kwarg_inputs, dict): + example_inputs = example_kwarg_inputs + else: + raise RuntimeError("example_kwarg_inputs should be a dict") + return trace_module( + func.__self__, + {"forward": example_inputs}, + None, + check_trace, + wrap_check_inputs(check_inputs), + check_tolerance, + strict, + _force_outplace, + _module_class, + example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict), + _store_inputs=_store_inputs, + ) + + # Special case for common case of passing a single Tensor + if ( + isinstance(example_inputs, (torch.Tensor, dict)) + and example_kwarg_inputs is None + ): + example_inputs = (example_inputs,) + # done primarily so that weird iterables fail here and not pybind11 code + elif example_kwarg_inputs is None and not isinstance(example_inputs, tuple): + example_inputs = tuple(example_inputs) + + var_lookup_fn = _create_interpreter_name_lookup_fn(0) + + if hasattr(func, "__self__") and isinstance(func.__self__, torch.nn.Module): + raise AttributeError( + "trace doesn't support compiling individual module's functions.\n" + "Please use trace_module" + ) + + name = _qualified_name(func) + if isinstance(example_kwarg_inputs, dict): + example_inputs = example_kwarg_inputs + traced = torch._C._create_function_from_trace_with_dict( + name, + func, + example_kwarg_inputs, + var_lookup_fn, + strict, + _force_outplace, + get_callable_argument_names(func), + ) + else: + traced = torch._C._create_function_from_trace( + name, + func, + example_inputs, + var_lookup_fn, + strict, + _force_outplace, + get_callable_argument_names(func), + ) + + # Check the trace against new traces created from user-specified inputs + if check_trace: + if check_inputs is not None: + _check_trace( + check_inputs, + func, + traced, + check_tolerance, + strict, + _force_outplace, + False, + _module_class, + example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict), + ) + else: + _check_trace( + [example_inputs], + func, + traced, + check_tolerance, + strict, + _force_outplace, + False, + _module_class, + example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict), + ) + + # Allow torch.compile() to inline + traced._torchdynamo_inline = func # type: ignore[attr-defined] + return traced + + +class _ExportType(str, Enum): + DIRECT_EXPORT = "DIRECT_EXPORT" + TRACE_AND_EXPORT = "TRACE_AND_EXPORT" + SOURCE_TO_SOURCE = "SOURCE_TO_SOURCE" + + def __str__(self) -> str: + return self.value + + +class _ExportOutcome(str, Enum): + SUCCESS = "SUCCESS" + FAILED_TO_EXPORT = "FAILED_TO_EXPORT" + FAILED_TO_RUN = "FAILED_TO_RUN" + ACCURACY_ERROR = "ACCURACY_ERROR" + + def __str__(self) -> str: + return self.value + + +def trace( + func, + example_inputs=None, + optimize=None, + check_trace=True, + check_inputs=None, + check_tolerance=1e-5, + strict=True, + _force_outplace=False, + _module_class=None, + _compilation_unit=_python_cu, + example_kwarg_inputs=None, + _store_inputs=True, +): + r""" + Trace a function and return an executable or :class:`ScriptFunction` that will be optimized using just-in-time compilation. + + Tracing is ideal for code that operates only on + ``Tensor``\\s and lists, dictionaries, and + tuples of ``Tensor``\\s. + + Using `torch.jit.trace` and `torch.jit.trace_module`, you can turn an + existing module or Python function into a TorchScript + :class:`ScriptFunction` or :class:`ScriptModule`. You must provide example + inputs, and we run the function, recording the operations performed on all + the tensors. + + * The resulting recording of a standalone function produces `ScriptFunction`. + * The resulting recording of `nn.Module.forward` or `nn.Module` produces + `ScriptModule`. + + This module also contains any parameters that the original + module had as well. + + Warning: + Tracing only correctly records functions and modules which are not data + dependent (e.g., do not have conditionals on data in tensors) and do not have + any untracked external dependencies (e.g., perform input/output or + access global variables). Tracing only records operations done when the given + function is run on the given tensors. Therefore, the returned + `ScriptModule` will always run the same traced graph on any input. This + has some important implications when your module is expected to run + different sets of operations, depending on the input and/or the module + state. For example, + + * Tracing will not record any control-flow like if-statements or loops. + When this control-flow is constant across your module, this is fine + and it often inlines the control-flow decisions. But sometimes the + control-flow is actually part of the model itself. For instance, a + recurrent network is a loop over the (possibly dynamic) length of an + input sequence. + * In the returned :class:`ScriptModule`, operations that have different + behaviors in ``training`` and ``eval`` modes will always behave as if + it is in the mode it was in during tracing, no matter which mode the + `ScriptModule` is in. + + In cases like these, tracing would not be appropriate and + :func:`scripting ` is a better choice. If you trace + such models, you may silently get incorrect results on subsequent + invocations of the model. The tracer will try to emit warnings when + doing something that may cause an incorrect trace to be produced. + + Args: + func (callable or torch.nn.Module): A Python function or `torch.nn.Module` + that will be run with `example_inputs`. `func` arguments and return + values must be tensors or (possibly nested) tuples that contain + tensors. When a module is passed `torch.jit.trace`, only the + ``forward`` method is run and traced (see :func:`torch.jit.trace + ` for details). + + Keyword arguments: + example_inputs (tuple or torch.Tensor or None, optional): A tuple of example + inputs that will be passed to the function while tracing. + Default: ``None``. Either this argument or ``example_kwarg_inputs`` + should be specified. The resulting trace can be run with inputs of + different types and shapes assuming the traced operations support those + types and shapes. `example_inputs` may also be a single Tensor in which + case it is automatically wrapped in a tuple. When the value is None, + ``example_kwarg_inputs`` should be specified. + + check_trace (``bool``, optional): Check if the same inputs run through + traced code produce the same outputs. Default: ``True``. You might want + to disable this if, for example, your network contains non- + deterministic ops or if you are sure that the network is correct despite + a checker failure. + + check_inputs (list of tuples, optional): A list of tuples of input + arguments that should be used to check the trace against what is + expected. Each tuple is equivalent to a set of input arguments that + would be specified in ``example_inputs``. For best results, pass in + a set of checking inputs representative of the space of shapes and + types of inputs you expect the network to see. If not specified, + the original ``example_inputs`` are used for checking + check_tolerance (float, optional): Floating-point comparison tolerance + to use in the checker procedure. This can be used to relax the + checker strictness in the event that results diverge numerically + for a known reason, such as operator fusion. + strict (``bool``, optional): run the tracer in a strict mode or not + (default: ``True``). Only turn this off when you want the tracer to + record your mutable container types (currently ``list``/``dict``) + and you are sure that the container you are using in your + problem is a ``constant`` structure and does not get used as + control flow (if, for) conditions. + example_kwarg_inputs (dict, optional): This parameter is a pack of keyword + arguments of example inputs that will be passed to the function while + tracing. Default: ``None``. Either this argument or ``example_inputs`` + should be specified. The dict will be unpacking by the arguments name + of the traced function. If the keys of the dict don't not match with + the traced function's arguments name, a runtime exception will be raised. + + Returns: + If `func` is `nn.Module` or ``forward`` of `nn.Module`, `trace` returns + a :class:`ScriptModule` object with a single ``forward`` method + containing the traced code. The returned `ScriptModule` will + have the same set of sub-modules and parameters as the original + ``nn.Module``. If ``func`` is a standalone function, ``trace`` + returns `ScriptFunction`. + + Example (tracing a function): + + .. testcode:: + + import torch + + def foo(x, y): + return 2 * x + y + + # Run `foo` with the provided inputs and record the tensor operations + traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3))) + + # `traced_foo` can now be run with the TorchScript interpreter or saved + # and loaded in a Python-free environment + + Example (tracing an existing module):: + + import torch + import torch.nn as nn + + class Net(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv2d(1, 1, 3) + + def forward(self, x): + return self.conv(x) + + n = Net() + example_weight = torch.rand(1, 1, 3, 3) + example_forward_input = torch.rand(1, 1, 3, 3) + + # Trace a specific method and construct `ScriptModule` with + # a single `forward` method + module = torch.jit.trace(n.forward, example_forward_input) + + # Trace a module (implicitly traces `forward`) and construct a + # `ScriptModule` with a single `forward` method + module = torch.jit.trace(n, example_forward_input) + + """ + if not _enabled: + return func + if optimize is not None: + warnings.warn( + "`optimize` is deprecated and has no effect. " + "Use `with torch.jit.optimized_execution()` instead", + FutureWarning, + stacklevel=2, + ) + + from torch._utils_internal import ( + check_if_torch_exportable, + log_torch_jit_trace_exportability, + log_torchscript_usage, + ) + + traced_func = _trace_impl( + func, + example_inputs, + optimize, + check_trace, + check_inputs, + check_tolerance, + strict, + _force_outplace, + _module_class, + _compilation_unit, + example_kwarg_inputs, + _store_inputs, + ) + log_torchscript_usage("trace", model_id=_get_model_id(traced_func)) + + if check_if_torch_exportable(): + from torch._export.converter import TS2EPConverter + from torch.export._trace import ( + _convert_ts_to_export_experimental, + _process_jit_trace_inputs_for_export, + ) + + traced_func_for_export = _trace_impl( + func, + example_inputs=example_inputs, + optimize=optimize, + check_trace=False, + check_inputs=check_inputs, + check_tolerance=check_tolerance, + strict=strict, + _force_outplace=_force_outplace, + _module_class=_module_class, + _compilation_unit=_compilation_unit, + example_kwarg_inputs=example_kwarg_inputs, + _store_inputs=_store_inputs, + ) + + export_args, _ = _process_jit_trace_inputs_for_export( + example_inputs, example_kwarg_inputs + ) + + def _log_exportability(func_to_export, export_func, export_args, export_type): + try: + traced_result = func_to_export(*export_args) + except Exception as e: + _ = e + log_torch_jit_trace_exportability( + "trace", str(export_type), str(_ExportOutcome.SUCCESS), "succeeded" + ) + return + + try: + ep_module = export_func(func_to_export, export_args) + except Exception as e: + log_torch_jit_trace_exportability( + "trace", + str(export_type), + str(_ExportOutcome.FAILED_TO_EXPORT), + str(e), + ) + return + + try: + export = ep_module(*export_args) + except Exception as e: + log_torch_jit_trace_exportability( + "trace", str(export_type), str(_ExportOutcome.FAILED_TO_RUN), str(e) + ) + return + + if not analyze_ts_result_with_export_result(export, traced_result): + log_torch_jit_trace_exportability( + "trace", + str(export_type), + str(_ExportOutcome.ACCURACY_ERROR), + "accuracy error", + ) + return + + log_torch_jit_trace_exportability( + "trace", str(export_type), str(_ExportOutcome.SUCCESS), "succeeded" + ) + + def _direct_export_and_lower(func, export_args): + return torch.export.export(func, export_args, strict=False).module() + + def _convert_ts_to_export_source_to_source(func, export_args): + return TS2EPConverter(func, export_args).convert().module() + + # torch.jit.trace is noop when the original module is torch.jit.ScriptModule + if not isinstance(traced_func_for_export, torch.jit.ScriptModule): + _log_exportability( + traced_func_for_export, + _direct_export_and_lower, + export_args, + _ExportType.DIRECT_EXPORT, + ) + + _log_exportability( + traced_func_for_export, + _convert_ts_to_export_experimental, + export_args, + _ExportType.TRACE_AND_EXPORT, + ) + _log_exportability( + traced_func_for_export, + _convert_ts_to_export_source_to_source, + export_args, + _ExportType.SOURCE_TO_SOURCE, + ) + + return traced_func + + +_trace_module_map: Optional[Dict[Any, Any]] = None + + +def trace_module( + mod, + inputs, + optimize=None, + check_trace=True, + check_inputs=None, + check_tolerance=1e-5, + strict=True, + _force_outplace=False, + _module_class=None, + _compilation_unit=_python_cu, + example_inputs_is_kwarg=False, + _store_inputs=True, +): + """ + Trace a module and return an executable :class:`ScriptModule` that will be optimized using just-in-time compilation. + + When a module is passed to :func:`torch.jit.trace `, only + the ``forward`` method is run and traced. With ``trace_module``, you can specify a dictionary of + method names to example inputs to trace (see the ``inputs``) argument below. + + See :func:`torch.jit.trace ` for more information on tracing. + + Args: + mod (torch.nn.Module): A ``torch.nn.Module`` containing methods whose names are + specified in ``inputs``. The given methods will be compiled + as a part of a single `ScriptModule`. + inputs (dict): A dict containing sample inputs indexed by method names in ``mod``. + The inputs will be passed to methods whose names correspond to inputs' + keys while tracing. + ``{ 'forward' : example_forward_input, 'method2': example_method2_input}`` + Keyword arguments: + check_trace (``bool``, optional): Check if the same inputs run through + traced code produce the same outputs. Default: ``True``. You might want + to disable this if, for example, your network contains non- + deterministic ops or if you are sure that the network is correct despite + a checker failure. + + check_inputs (list of dicts, optional): A list of dicts of input arguments that should be used + to check the trace against what is expected. Each tuple + is equivalent to a set of input arguments that would + be specified in ``inputs``. For best results, pass in a + set of checking inputs representative of the space of + shapes and types of inputs you expect the network to see. + If not specified, the original ``inputs`` are used for checking + check_tolerance (float, optional): Floating-point comparison tolerance to use in the checker procedure. + This can be used to relax the checker strictness in the event that + results diverge numerically for a known reason, such as operator fusion. + example_inputs_is_kwarg (``bool``, optional): This parameter indicate whether the example inputs is a pack + pack of keyword arguments. Default: ``False``. + + Returns: + A :class:`ScriptModule` object with a single ``forward`` method containing the traced code. + When ``func`` is a ``torch.nn.Module``, the returned :class:`ScriptModule` will have the same set of + sub-modules and parameters as ``func``. + + Example (tracing a module with multiple methods):: + + import torch + import torch.nn as nn + + class Net(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv2d(1, 1, 3) + + def forward(self, x): + return self.conv(x) + + def weighted_kernel_sum(self, weight): + return weight * self.conv.weight + + + n = Net() + example_weight = torch.rand(1, 1, 3, 3) + example_forward_input = torch.rand(1, 1, 3, 3) + + # Trace a specific method and construct `ScriptModule` with + # a single `forward` method + module = torch.jit.trace(n.forward, example_forward_input) + + # Trace a module (implicitly traces `forward`) and construct a + # `ScriptModule` with a single `forward` method + module = torch.jit.trace(n, example_forward_input) + + # Trace specific methods on a module (specified in `inputs`), constructs + # a `ScriptModule` with `forward` and `weighted_kernel_sum` methods + inputs = {'forward' : example_forward_input, 'weighted_kernel_sum' : example_weight} + module = torch.jit.trace_module(n, inputs) + + """ + if not _enabled: + return mod + if optimize is not None: + warnings.warn( + "`optimize` is deprecated and has no effect. " + "Use `with torch.jit.optimized_execution()` instead", + FutureWarning, + stacklevel=2, + ) + + var_lookup_fn = _create_interpreter_name_lookup_fn(0) + + if not isinstance(mod, torch.nn.Module): + raise AttributeError("expected torch.nn.Module as the first argument") + + if not isinstance(inputs, dict): + raise AttributeError("expected a dictionary of (method_name, input) pairs") + + old_module_map = torch.jit._trace._trace_module_map + try: + trace_module_map: Dict[Any, Any] = {} + + def register_submods(mod, prefix): + for name, child in mod.named_children(): + submod_qualname = prefix + "." + name + trace_module_map[child] = submod_qualname + register_submods(child, submod_qualname) + + trace_module_map["__module"] = mod + torch.jit._trace._trace_module_map = trace_module_map + register_submods(mod, "__module") + + module = make_module(mod, _module_class, _compilation_unit) + + for method_name, example_inputs in inputs.items(): + if method_name == "forward": + # "forward" is a special case because we need to trace + # `Module.__call__`, which sets up some extra tracing, but uses + # argument names of the real `Module.forward` method. + func = mod + forward_method = getattr(mod, method_name) + argument_names = get_callable_argument_names(forward_method) + else: + func = getattr(mod, method_name) + argument_names = get_callable_argument_names(func) + + if isinstance(example_inputs, dict) and example_inputs_is_kwarg: + # Raise exception when the user provided key names are not aligned with forward() method's arguments' name/ + for key in example_inputs: + if key not in argument_names: + valid_arguments = "[" + ",".join(argument_names) + "]" + raise NameError( + f"""'{key}' is not in forward() method's arguments, + valid arguments name are {valid_arguments}""" + ) + module._c._create_method_from_trace_with_dict( + method_name, + func, + example_inputs, + var_lookup_fn, + strict, + _force_outplace, + argument_names, + _store_inputs, + ) + else: + example_inputs = make_tuple(example_inputs) + module._c._create_method_from_trace( + method_name, + func, + example_inputs, + var_lookup_fn, + strict, + _force_outplace, + argument_names, + _store_inputs, + ) + + check_trace_method = module._c._get_method(method_name) + + # Check the trace against new traces created from user-specified inputs + if check_trace: + if check_inputs is not None: + _check_trace( + check_inputs, + func, + check_trace_method, + check_tolerance, + strict, + _force_outplace, + True, + _module_class, + example_inputs_is_kwarg=example_inputs_is_kwarg, + ) + else: + _check_trace( + [inputs], + func, + check_trace_method, + check_tolerance, + strict, + _force_outplace, + True, + _module_class, + example_inputs_is_kwarg=example_inputs_is_kwarg, + ) + finally: + torch.jit._trace._trace_module_map = old_module_map + + return module + + +def is_tracing(): + """Return a boolean value. + + Returns ``True`` in tracing (if a function is called during the + tracing of code with ``torch.jit.trace``) and ``False`` otherwise. + """ + if is_scripting(): + return False + return torch._C._is_tracing() + + +class TracedModule(ScriptModule): + _disable_script_meta = True + + def __init__(self, orig, id_set=None, _compilation_unit=None): + # XXX: orig can be a nn.Module or a function! + super().__init__() + assert isinstance(orig, torch.nn.Module) + + # Copy a subset of `orig` to a temporary nn.Module. + # This is a way to customize what will actually get compiled by create_script_module + id_set = set() + + # This allows us to preserve the original module's qualified name by defining a new + # type with the attribute _jit_override_qualname. In torch._jit_internal._qualified_name + # we have a special case that will look up this attribute to override whatever qualname + # we would get from the python type system + class QualnameWrapper(torch.nn.Module): + pass + + QualnameWrapper._jit_override_qualname = torch._jit_internal._qualified_name( # type: ignore[attr-defined] + type(orig) + ) + + tmp_module = QualnameWrapper() + + def check_unique(param): + if param in id_set: + raise ValueError( + "TracedModules don't support parameter sharing between modules" + ) + id_set.add(param) + + tmp_module.training = orig.training + + for name, param in orig._parameters.items(): + if param is not None: + tmp_module._parameters[name] = param + check_unique(param) + for name, buf in orig._buffers.items(): + if buf is not None: + tmp_module._buffers[name] = buf + check_unique(buf) + for name, val in orig.__dict__.items(): + if ( + torch._C._jit_is_script_object(val) + and name not in orig._parameters + and name not in orig._buffers + ): + setattr(tmp_module, name, val) + + if orig._backward_hooks: + raise ValueError( + "Modules that have backward hooks assigned can't be compiled: " + + str(orig) + ) + + for name, submodule in orig._modules.items(): + if submodule is None: + continue + tmp_module._modules[name] = make_module( + submodule, TracedModule, _compilation_unit=None + ) + + script_module = torch.jit._recursive.create_script_module( + tmp_module, lambda module: (), share_types=False, is_tracing=True + ) + + self.__dict__["_name"] = type(orig).__name__ + self.__dict__["_actual_script_module"] = script_module + for name in ("_parameters", "_buffers", "_modules", "training"): + delattr(self, name) + + def forward(self, *args, **kwargs): + raise RuntimeError("Trace submodules cannot be called.") + + def __getattr__(self, attr): + if "_actual_script_module" not in self.__dict__: + return super().__getattr__(attr) + return getattr(self._actual_script_module, attr) + + def __setattr__(self, attr, value): + if "_actual_script_module" not in self.__dict__: + return super().__setattr__(attr, value) + setattr(self._actual_script_module, attr, value) + + def _get_name(self): + return self._name + + def extra_repr(self): + return f"original_name={self._name}" + + +class TopLevelTracedModule(TracedModule): + forward: Callable[..., Any] = _CachedForward() # type: ignore[assignment] + + def _reconstruct(self, cpp_module): + """ + Re-construct an instance of TopLevelTracedModule using an instance of a C++ module. + + Args: + cpp_module: The C++ module that this TopLevelTracedModule will be rebuilt around. + """ + self.__dict__["_actual_script_module"]._reconstruct(cpp_module) + + +def _script_if_tracing(fn: Callable[P, R]) -> Callable[P, R]: + @functools.wraps(fn) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + if not is_tracing(): + # Not tracing, don't do anything + return fn(*args, **kwargs) + + compiled_fn: Callable[P, R] = script(wrapper.__original_fn) # type: ignore[attr-defined] + return compiled_fn(*args, **kwargs) + + wrapper.__original_fn = fn # type: ignore[attr-defined] + wrapper.__script_if_tracing_wrapper = True # type: ignore[attr-defined] + + return wrapper + + +def _get_trace_graph( + f, + args=(), + kwargs=None, + strict=True, + _force_outplace=False, + return_inputs=False, + _return_inputs_states=False, +): + """Return a tuple on tracing a function or model. + + .. warning:: + This function is internal-only and should only be used by the ONNX + exporter. If you are trying to get a graph through tracing, please go + through the public API instead:: + + trace = torch.jit.trace(nn.LSTMCell(), (input, hidden)) + trace_graph = trace.graph + + Trace a function or model, returning a tuple consisting of the both the + *trace* of an execution, as well as the original return value. If return_inputs, + also returns the trace inputs as part of the tuple + + Tracing is guaranteed not to change the semantics of the function/module + that is traced. + + Args: + f (torch.nn.Module or function): the function or module + to be traced. + args (tuple or Tensor): the positional arguments to pass to the + function/module to be traced. A non-tuple is assumed to + be a single positional argument to be passed to the model. + kwargs (dict): the keyword arguments to pass to the function/module + to be traced. + + Example (trace a cell): + + .. testcode:: + + trace = torch.jit.trace(nn.LSTMCell(), (input, hidden)) + """ + if kwargs is None: + kwargs = {} + if not isinstance(args, tuple): + args = (args,) + outs = ONNXTracedModule( + f, strict, _force_outplace, return_inputs, _return_inputs_states + )(*args, **kwargs) + return outs diff --git a/lib/python3.12/site-packages/torch/jit/annotations.py b/lib/python3.12/site-packages/torch/jit/annotations.py new file mode 100644 index 0000000000000000000000000000000000000000..0337342edaf97be421b037dafdde199a3d29af42 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/annotations.py @@ -0,0 +1,551 @@ +# mypy: allow-untyped-defs +import ast +import builtins +import dis +import enum +import inspect +import re +import typing +import warnings +from textwrap import dedent +from typing import Type + +import torch +from torch._C import ( + _GeneratorType, + AnyType, + AwaitType, + BoolType, + ComplexType, + DeviceObjType, + DictType, + EnumType, + FloatType, + FutureType, + InterfaceType, + IntType, + ListType, + NoneType, + NumberType, + OptionalType, + StreamObjType, + StringType, + TensorType, + TupleType, + UnionType, +) +from torch._jit_internal import ( # type: ignore[attr-defined] + _Await, + _qualified_name, + Any, + BroadcastingList1, + BroadcastingList2, + BroadcastingList3, + Dict, + Future, + is_await, + is_dict, + is_future, + is_ignored_fn, + is_list, + is_optional, + is_tuple, + is_union, + List, + Optional, + Tuple, + Union, +) +from torch._sources import get_source_lines_and_file + +from ._state import _get_script_class + + +if torch.distributed.rpc.is_available(): + from torch._C import RRefType + from torch._jit_internal import is_rref, RRef + +from torch._ops import OpOverloadPacket + + +class Module: + def __init__(self, name, members): + self.name = name + self.members = members + + def __getattr__(self, name): + try: + return self.members[name] + except KeyError: + raise RuntimeError( + f"Module {self.name} has no member called {name}" + ) from None + + +class EvalEnv: + env = { + "torch": Module("torch", {"Tensor": torch.Tensor}), + "Tensor": torch.Tensor, + "typing": Module("typing", {"Tuple": Tuple}), + "Tuple": Tuple, + "List": List, + "Dict": Dict, + "Optional": Optional, + "Union": Union, + "Future": Future, + "Await": _Await, + } + + def __init__(self, rcb): + self.rcb = rcb + if torch.distributed.rpc.is_available(): + self.env["RRef"] = RRef + + def __getitem__(self, name): + if name in self.env: + return self.env[name] + if self.rcb is not None: + return self.rcb(name) + return getattr(builtins, name, None) + + +def get_signature(fn, rcb, loc, is_method): + if isinstance(fn, OpOverloadPacket): + signature = try_real_annotations(fn.op, loc) + else: + signature = try_real_annotations(fn, loc) + if signature is not None and is_method: + # If this is a method, then the signature will include a type for + # `self`, but type comments do not contain a `self`. So strip it + # away here so everything is consistent (`inspect.ismethod` does + # not work here since `fn` is unbound at this point) + param_types, return_type = signature + param_types = param_types[1:] + signature = (param_types, return_type) + + if signature is None: + type_line, source = None, None + try: + source = dedent("".join(get_source_lines_and_file(fn)[0])) + type_line = get_type_line(source) + except TypeError: + pass + # This might happen both because we failed to get the source of fn, or + # because it didn't have any annotations. + if type_line is not None: + signature = parse_type_line(type_line, rcb, loc) + + return signature + + +def is_function_or_method(the_callable): + # A stricter version of `inspect.isroutine` that does not pass for built-in + # functions + return inspect.isfunction(the_callable) or inspect.ismethod(the_callable) + + +def is_vararg(the_callable): + if not is_function_or_method(the_callable) and callable(the_callable): # noqa: B004 + # If `the_callable` is a class, de-sugar the call so we can still get + # the signature + the_callable = the_callable.__call__ + + if is_function_or_method(the_callable): + return inspect.getfullargspec(the_callable).varargs is not None + else: + return False + + +def get_param_names(fn, n_args): + if isinstance(fn, OpOverloadPacket): + fn = fn.op + + if ( + not is_function_or_method(fn) + and callable(fn) + and is_function_or_method(fn.__call__) + ): # noqa: B004 + # De-sugar calls to classes + fn = fn.__call__ + + if is_function_or_method(fn): + if is_ignored_fn(fn): + fn = inspect.unwrap(fn) + return inspect.getfullargspec(fn).args + else: + # The `fn` was not a method or function (maybe a class with a __call__ + # method, so use a default param name list) + return [str(i) for i in range(n_args)] + + +def check_fn(fn, loc): + # Make sure the function definition is not a class instantiation + try: + source = dedent("".join(get_source_lines_and_file(fn)[0])) + except (OSError, TypeError): + return + if source is None: + return + + py_ast = ast.parse(source) + if len(py_ast.body) == 1 and isinstance(py_ast.body[0], ast.ClassDef): + raise torch.jit.frontend.FrontendError( + loc, + f"Cannot instantiate class '{py_ast.body[0].name}' in a script function", + ) + if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef): + raise torch.jit.frontend.FrontendError( + loc, "Expected a single top-level function" + ) + + +def _eval_no_call(stmt, glob, loc): + """Evaluate statement as long as it does not contain any method/function calls.""" + bytecode = compile(stmt, "", mode="eval") + for insn in dis.get_instructions(bytecode): + if "CALL" in insn.opname: + raise RuntimeError( + f"Type annotation should not contain calls, but '{stmt}' does" + ) + return eval(bytecode, glob, loc) # type: ignore[arg-type] # noqa: P204 + + +def parse_type_line(type_line, rcb, loc): + """Parse a type annotation specified as a comment. + + Example inputs: + # type: (Tensor, torch.Tensor) -> Tuple[Tensor] + # type: (Tensor, Tuple[Tensor, Tensor]) -> Tensor + """ + arg_ann_str, ret_ann_str = split_type_line(type_line) + + try: + arg_ann = _eval_no_call(arg_ann_str, {}, EvalEnv(rcb)) + except (NameError, SyntaxError) as e: + raise RuntimeError( + "Failed to parse the argument list of a type annotation" + ) from e + + if not isinstance(arg_ann, tuple): + arg_ann = (arg_ann,) + + try: + ret_ann = _eval_no_call(ret_ann_str, {}, EvalEnv(rcb)) + except (NameError, SyntaxError) as e: + raise RuntimeError( + "Failed to parse the return type of a type annotation" + ) from e + + arg_types = [ann_to_type(ann, loc) for ann in arg_ann] + return arg_types, ann_to_type(ret_ann, loc) + + +def get_type_line(source): + """Try to find the line containing a comment with the type annotation.""" + type_comment = "# type:" + + lines = source.split("\n") + lines = list(enumerate(lines)) + type_lines = list(filter(lambda line: type_comment in line[1], lines)) + # `type: ignore` comments may be needed in JIT'ed functions for mypy, due + # to the hack in torch/_VF.py. + + # An ignore type comment can be of following format: + # 1) type: ignore + # 2) type: ignore[rule-code] + # This ignore statement must be at the end of the line + + # adding an extra backslash before the space, to avoid triggering + # one of the checks in .github/workflows/lint.yml + type_pattern = re.compile("# type:\\ ignore(\\[[a-zA-Z-]+\\])?$") + type_lines = list(filter(lambda line: not type_pattern.search(line[1]), type_lines)) + + if len(type_lines) == 0: + # Catch common typo patterns like extra spaces, typo in 'ignore', etc. + wrong_type_pattern = re.compile("#[\t ]*type[\t ]*(?!: ignore(\\[.*\\])?$):") + wrong_type_lines = list( + filter(lambda line: wrong_type_pattern.search(line[1]), lines) + ) + if len(wrong_type_lines) > 0: + raise RuntimeError( + "The annotation prefix in line " + + str(wrong_type_lines[0][0]) + + " is probably invalid.\nIt must be '# type:'" + + "\nSee PEP 484 (https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)" # noqa: B950 + + "\nfor examples" + ) + return None + elif len(type_lines) == 1: + # Only 1 type line, quit now + return type_lines[0][1].strip() + + # Parse split up argument types according to PEP 484 + # https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code + return_line = None + parameter_type_lines = [] + for line_num, line in type_lines: + if "# type: (...) -> " in line: + return_line = (line_num, line) + break + elif type_comment in line: + parameter_type_lines.append(line) + if return_line is None: + raise RuntimeError( + "Return type line '# type: (...) -> ...' not found on multiline " + "type annotation\nfor type lines:\n" + + "\n".join([line[1] for line in type_lines]) + + "\n(See PEP 484 https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)" + ) + + def get_parameter_type(line): + item_type = line[line.find(type_comment) + len(type_comment) :] + return item_type.strip() + + types = map(get_parameter_type, parameter_type_lines) + parameter_types = ", ".join(types) + + return return_line[1].replace("...", parameter_types) + + +def split_type_line(type_line): + """Split the comment with the type annotation into parts for argument and return types. + + For example, for an input of: + # type: (Tensor, torch.Tensor) -> Tuple[Tensor, Tensor] + + This function will return: + ("(Tensor, torch.Tensor)", "Tuple[Tensor, Tensor]") + + """ + start_offset = len("# type:") + try: + arrow_pos = type_line.index("->") + except ValueError: + raise RuntimeError( + "Syntax error in type annotation (couldn't find `->`)" + ) from None + return type_line[start_offset:arrow_pos].strip(), type_line[arrow_pos + 2 :].strip() + + +def try_real_annotations(fn, loc): + """Try to use the Py3.5+ annotation syntax to get the type.""" + try: + # Note: anything annotated as `Optional[T]` will automatically + # be returned as `Union[T, None]` per + # https://github.com/python/typing/blob/master/src/typing.py#L850 + sig = inspect.signature(fn) + except ValueError: + return None + + all_annots = [sig.return_annotation] + [ + p.annotation for p in sig.parameters.values() + ] + if all(ann is sig.empty for ann in all_annots): + return None + + arg_types = [ann_to_type(p.annotation, loc) for p in sig.parameters.values()] + return_type = ann_to_type(sig.return_annotation, loc) + return arg_types, return_type + + +# Finds common type for enum values belonging to an Enum class. If not all +# values have the same type, AnyType is returned. +def get_enum_value_type(e: Type[enum.Enum], loc): + enum_values: List[enum.Enum] = list(e) + if not enum_values: + raise ValueError(f"No enum values defined for: '{e.__class__}'") + + types = {type(v.value) for v in enum_values} + ir_types = [try_ann_to_type(t, loc) for t in types] + + # If Enum values are of different types, an exception will be raised here. + # Even though Python supports this case, we chose to not implement it to + # avoid overcomplicate logic here for a rare use case. Please report a + # feature request if you find it necessary. + res = torch._C.unify_type_list(ir_types) + if not res: + return AnyType.get() + return res + + +def is_tensor(ann): + if issubclass(ann, torch.Tensor): + return True + + if issubclass( + ann, + ( + torch.LongTensor, + torch.DoubleTensor, + torch.FloatTensor, + torch.IntTensor, + torch.ShortTensor, + torch.HalfTensor, + torch.CharTensor, + torch.ByteTensor, + torch.BoolTensor, + ), + ): + warnings.warn( + "TorchScript will treat type annotations of Tensor " + "dtype-specific subtypes as if they are normal Tensors. " + "dtype constraints are not enforced in compilation either." + ) + return True + + return False + + +def _fake_rcb(inp): + return None + + +def try_ann_to_type(ann, loc, rcb=None): + ann_args = typing.get_args(ann) # always returns a tuple! + + if ann is inspect.Signature.empty: + return TensorType.getInferred() + if ann is None: + return NoneType.get() + if inspect.isclass(ann) and is_tensor(ann): + return TensorType.get() + if is_tuple(ann): + # Special case for the empty Tuple type annotation `Tuple[()]` + if len(ann_args) == 1 and ann_args[0] == (): + return TupleType([]) + return TupleType([try_ann_to_type(a, loc) for a in ann_args]) + if is_list(ann): + elem_type = try_ann_to_type(ann_args[0], loc) + if elem_type: + return ListType(elem_type) + if is_dict(ann): + key = try_ann_to_type(ann_args[0], loc) + value = try_ann_to_type(ann_args[1], loc) + # Raise error if key or value is None + if key is None: + raise ValueError( + f"Unknown type annotation: '{ann_args[0]}' at {loc.highlight()}" + ) + if value is None: + raise ValueError( + f"Unknown type annotation: '{ann_args[1]}' at {loc.highlight()}" + ) + return DictType(key, value) + if is_optional(ann): + if issubclass(ann_args[1], type(None)): + contained = ann_args[0] + else: + contained = ann_args[1] + valid_type = try_ann_to_type(contained, loc) + msg = "Unsupported annotation {} could not be resolved because {} could not be resolved. At\n{}" + assert valid_type, msg.format(repr(ann), repr(contained), repr(loc)) + return OptionalType(valid_type) + if is_union(ann): + # TODO: this is hack to recognize NumberType + if set(ann_args) == {int, float, complex}: + return NumberType.get() + inner: List = [] + # We need these extra checks because both `None` and invalid + # values will return `None` + # TODO: Determine if the other cases need to be fixed as well + for a in typing.get_args(ann): + if a is None: + inner.append(NoneType.get()) + maybe_type = try_ann_to_type(a, loc) + msg = "Unsupported annotation {} could not be resolved because {} could not be resolved. At\n{}" + assert maybe_type, msg.format(repr(ann), repr(maybe_type), repr(loc)) + inner.append(maybe_type) + return UnionType(inner) # type: ignore[arg-type] + if torch.distributed.rpc.is_available() and is_rref(ann): + return RRefType(try_ann_to_type(ann_args[0], loc)) + if is_future(ann): + return FutureType(try_ann_to_type(ann_args[0], loc)) + if is_await(ann): + elementType = try_ann_to_type(ann_args[0], loc) if ann_args else AnyType.get() + return AwaitType(elementType) + if ann is float: + return FloatType.get() + if ann is complex: + return ComplexType.get() + if ann is int or ann is torch.SymInt: + return IntType.get() + if ann is str: + return StringType.get() + if ann is bool: + return BoolType.get() + if ann is Any: + return AnyType.get() + if ann is type(None): + return NoneType.get() + if inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"): + return InterfaceType(ann.__torch_script_interface__) + if ann is torch.device: + return DeviceObjType.get() + if ann is torch.Generator: + return _GeneratorType.get() + if ann is torch.Stream: + return StreamObjType.get() + if ann is torch.dtype: + return IntType.get() # dtype not yet bound in as its own type + if inspect.isclass(ann) and issubclass(ann, enum.Enum): + if _get_script_class(ann) is None: + scripted_class = torch.jit._script._recursive_compile_class(ann, loc) + name = scripted_class.qualified_name() + else: + name = _qualified_name(ann) + return EnumType(name, get_enum_value_type(ann, loc), list(ann)) + if inspect.isclass(ann): + maybe_script_class = _get_script_class(ann) + if maybe_script_class is not None: + return maybe_script_class + if torch._jit_internal.can_compile_class(ann): + return torch.jit._script._recursive_compile_class(ann, loc) + + # Maybe resolve a NamedTuple to a Tuple Type + if rcb is None: + rcb = _fake_rcb + return torch._C._resolve_type_from_object(ann, loc, rcb) + + +def ann_to_type(ann, loc, rcb=None): + the_type = try_ann_to_type(ann, loc, rcb) + if the_type is not None: + return the_type + raise ValueError(f"Unknown type annotation: '{ann}' at {loc.highlight()}") + + +__all__ = [ + "Any", + "List", + "BroadcastingList1", + "BroadcastingList2", + "BroadcastingList3", + "Tuple", + "is_tuple", + "is_list", + "Dict", + "is_dict", + "is_optional", + "is_union", + "TensorType", + "TupleType", + "FloatType", + "ComplexType", + "IntType", + "ListType", + "StringType", + "DictType", + "AnyType", + "Module", + # TODO: Consider not exporting these during wildcard import (reserve + # that for the types; for idiomatic typing code.) + "get_signature", + "check_fn", + "get_param_names", + "parse_type_line", + "get_type_line", + "split_type_line", + "try_real_annotations", + "try_ann_to_type", + "ann_to_type", +] diff --git a/lib/python3.12/site-packages/torch/jit/frontend.py b/lib/python3.12/site-packages/torch/jit/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..fd7c8f22e288ac8b6298133b36f96807b86ba05a --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/frontend.py @@ -0,0 +1,1286 @@ +# mypy: allow-untyped-defs +import ast +import dataclasses +import inspect +import re +import string +import sys +from collections import namedtuple +from textwrap import dedent +from typing import List, Tuple # noqa: F401 + +import torch +import torch.jit.annotations +from torch import _jit_internal +from torch._C._jit_tree_views import ( + Apply, + Assert, + Assign, + Attribute, + AugAssign, + BinOp, + Break, + ClassDef, + Const, + Continue, + Decl, + Def, + Delete, + DictComp, + DictLiteral, + Dots, + EmptyTypeAnnotation, + ExprStmt, + FalseLiteral, + For, + Ident, + If, + ListComp, + ListLiteral, + NoneLiteral, + Param, + Pass, + Property, + Raise, + Return, + Select, + SliceExpr, + Starred, + Stmt, + StringLiteral, + Subscript, + TernaryIf, + TrueLiteral, + TupleLiteral, + UnaryOp, + Var, + While, + With, + WithItem, +) +from torch._jit_internal import ( # noqa: F401 + _is_drop_fn, + FunctionModifiers, + is_static_fn, + should_drop, +) +from torch._sources import ( + get_source_lines_and_file, + make_source_context, + parse_def, + ParsedDef as _ParsedDef, +) +from torch.jit._dataclass_impls import DATACLASS_MAGIC_METHODS +from torch.jit._monkeytype_config import get_qualified_name, monkeytype_trace + + +_IS_ASTUNPARSE_INSTALLED = False +try: + import astunparse # type: ignore[import] + + _IS_ASTUNPARSE_INSTALLED = True +except ImportError: + pass + +# Borrowed from cPython implementation +# https://github.com/python/cpython/blob/561612d8456cfab5672c9b445521113b847bd6b3/Lib/textwrap.py#L411# + +_reserved_prefix = "__jit" +_reserved_names = {"print"} +_identifier_chars = set(string.ascii_lowercase + string.ascii_uppercase + string.digits) + + +def is_reserved_name(name): + return name.startswith(_reserved_prefix) or name in _reserved_names + + +pretty_node_names = { + ast.FunctionDef: "function definitions", + ast.For: "for loops", + ast.Delete: "del statements", + ast.ClassDef: "class definitions", + ast.With: "with statements", + ast.Raise: "raise statements", + ast.Assert: "assertions", + ast.Import: "import statements", + ast.ImportFrom: "import statements", + ast.Global: "global variables", + ast.Break: "break statements", + ast.Continue: "continue statements", +} + +node_start_tokens = { + ast.FunctionDef: "def", + ast.For: "for", + ast.Delete: "del", + ast.ClassDef: "class", + ast.With: "with", + ast.Raise: "raise", + ast.Assert: "assert", + ast.Import: "import", + ast.ImportFrom: "from", + ast.Global: "global", + ast.Break: "break", + ast.Continue: "continue", +} + +pretty_node_names.update( + { + ast.AsyncFunctionDef: "async function definitions", + ast.AsyncFor: "async for loops", + ast.AsyncWith: "async with statements", + ast.Try: "try blocks", + ast.Nonlocal: "nonlocal variables", + } +) + +node_start_tokens.update( + { + ast.AsyncFunctionDef: "async def", + ast.AsyncFor: "async for", + ast.AsyncWith: "async with", + ast.Try: "try", + ast.Nonlocal: "nonlocal", + } +) + +pretty_node_names.update( + { + ast.AnnAssign: "annotated assignments", + } +) +# NB: no specific token for AnnAssign + + +class FrontendError(Exception): + def __init__(self, source_range, msg): + self.source_range = source_range + self.msg = msg + + # This has to be instantiated here so the ErrorReport is accurate to the + # call stack when the FrontendError was raised + self.error_report = torch._C.ErrorReport(self.source_range) + + def __str__(self): + return self.msg + self.error_report.what().lstrip() + + +class NotSupportedError(FrontendError): + pass + + +class UnsupportedNodeError(NotSupportedError): + def __init__(self, ctx, offending_node, reason=""): + # If we don't have a specific token, we default to length of 1 + node_type = type(offending_node) + range_len = len(node_start_tokens.get(node_type, " ")) + source_range = ctx.make_range( + offending_node.lineno, + offending_node.col_offset, + offending_node.col_offset + range_len, + ) + feature_name = pretty_node_names.get(node_type, node_type.__name__) + msg = f"{feature_name} {reason + ' ' if reason else ''}aren't supported" + super().__init__(source_range, msg) + + +class FrontendTypeError(FrontendError): + pass + + +def build_withitems(ctx, items): + items = [build_withitem(ctx, i) for i in items] + return list(items) + + +def build_stmts(ctx, stmts): + stmts = [build_stmt(ctx, s) for s in stmts] + return list(filter(None, stmts)) + + +def get_class_properties(cls, self_name): + """ + Get a list of Property objects representing the properties of a class. + + Args: + cls: The class to get properties of. + self_name: The name of the class that the properties should belong to. + Returns: + A list of Property objects corresponding to the properties of cls. Property + here refers to the subclass of TreeView. + """ + props = inspect.getmembers(cls, predicate=lambda m: isinstance(m, property)) + # Any property that should not compiled must be in this list on the Module. + unused_properties = getattr(cls, "__jit_unused_properties__", []) + + # Create Property TreeView objects from inspected property objects. + properties = [] + for prop in props: + if prop[0] not in unused_properties and not should_drop(prop[1].fget): + getter = get_jit_def( + prop[1].fget, f"__{prop[0]}_getter", self_name=self_name + ) + setter = ( + get_jit_def(prop[1].fset, f"__{prop[0]}_setter", self_name=self_name) + if prop[1].fset + else None + ) + properties.append( + Property(getter.range(), Ident(getter.range(), prop[0]), getter, setter) + ) + + return properties + + +def get_class_assigns(ctx, cls_ast): + assigns = [] + + def maybe_build_assign(builder, entry): + nonlocal assigns + try: + assigns.append(builder(ctx, entry)) + except NotSupportedError: + pass + + for entry in cls_ast.body: + if isinstance(entry, ast.Assign): + maybe_build_assign(StmtBuilder.build_Assign, entry) + elif isinstance(entry, ast.AnnAssign): + maybe_build_assign(StmtBuilder.build_AnnAssign, entry) + return assigns + + +def get_jit_class_def(cls, self_name): + """Get definitions for each method within the current class independently. + + Args: + cls: The class to get definition of. + self_name: The name of the class that the properties should belong to. + + Returns: + torch._C._jit_tree_views.ClassDef: A representation of the class, + the methods in the class and their definition as a tree. + """ + # TODO: proper overriding analysis when implementing class inheritance + methods = inspect.getmembers( + cls, + predicate=lambda m: (inspect.ismethod(m) or inspect.isfunction(m)) + and not is_static_fn(cls, m.__name__) + and m.__name__ in cls.__dict__ + and not _is_drop_fn(m), + ) + + def is_classmethod(fn): + return inspect.ismethod(fn) and getattr(fn, "__self__", None) == cls + + # Get and parse the source code for this class + sourcelines, file_lineno, filename = get_source_lines_and_file( + cls, torch._C.ErrorReport.call_stack() + ) + source = "".join(sourcelines) + + dedent_src = dedent(source) + py_ast = ast.parse(dedent_src) + + class_ast = py_ast.body[0] + assert isinstance(class_ast, ast.ClassDef) + + # Special case for dataclasses. In general we need access to the source code for + # an object in order to JIT compile it. But the dataclasses module dynamically synthesizes + # magic methods for classes, and we can't get the source code for these methods. As a + # workaround, we synthesize TorchScript-friendly implementations ourselves. + if dataclasses.is_dataclass(cls): + # Detect whether the user manually implemented any of the magic methods. If they did, + # we don't want to synthesize/override them. + overrides = { + method.name + for method in class_ast.body + if isinstance(method, ast.FunctionDef) + and method.name in DATACLASS_MAGIC_METHODS + } + for i, (name, _) in enumerate(methods): + # Is this a magic method we can synthesize? + synthesizer_fn = DATACLASS_MAGIC_METHODS.get(name) + if synthesizer_fn and name not in overrides: + parsed_def = synthesizer_fn(cls) + methods[i] = name, parsed_def + func = getattr(cls, name) + _jit_internal.loader.cache(func, parsed_def.source) + + method_defs = [ + get_jit_def(obj, name, self_name=self_name, is_classmethod=is_classmethod(obj)) + for (name, obj) in methods + ] + properties = get_class_properties(cls, self_name) + + leading_whitespace_len = len(source.split("\n", 1)[0]) - len( + dedent_src.split("\n", 1)[0] + ) + ctx = make_source_context( + source, filename, file_lineno, leading_whitespace_len, False + ) + assigns = get_class_assigns(ctx, class_ast) + + return build_class_def(ctx, class_ast, method_defs, properties, self_name, assigns) + + +def get_jit_def(fn, def_name, self_name=None, is_classmethod=False): + """ + Build a JIT AST (TreeView) from the given function. + + Args: + fn: A function object to compile or a pre-parsed ParsedDef object + def_name: The name to give to the resulting AST object. This is not + always the same as `fn.__name__`, for example: + def _forward(self): + ... + forward = _forward + In this case, the `__name__` attribute of the function object is "_forward", + but we want the result AST to have the name "forward". + self_name: If this function is a method, what the type name of `self` is. + """ + parsed_def = parse_def(fn) if not isinstance(fn, _ParsedDef) else fn + type_line = torch.jit.annotations.get_type_line(parsed_def.source) + fn_def = parsed_def.ast.body[0] + + if is_classmethod: + arg_name = fn_def.args.args[0].arg # type:ignore[union-attr] + # Insert a statement that assigns the first argument to the class + assign_stmt = ast.parse(f"{arg_name} = {self_name}").body[0] + fn_def.body.insert(0, assign_stmt) # type:ignore[union-attr] + + # Swap out the function signature and body if it is unused + if should_drop(fn): + unused_fn_def = ast.parse( + 'def unused_fn(self: Any):\n\traise RuntimeError("Cannot call @unused methods")' + ) + if len(unused_fn_def.body) != 1 or not isinstance( + unused_fn_def.body[0], ast.FunctionDef + ): + raise RuntimeError( + f"Expected a single top-level function: {parsed_def.filename}:{parsed_def.file_lineno}" + ) + unused_def = unused_fn_def.body[0] + fn_def.body = unused_def.body # type:ignore[union-attr] + # kwarg/vararg not supported by `build_def` + fn_def.args.kwarg = fn_def.args.vararg = None # type:ignore[union-attr] + for arg in fn_def.args.args + fn_def.args.kwonlyargs: # type:ignore[union-attr] + # Replace potentially unsupported type annotations by "Any" + arg.annotation = unused_def.args.args[0].annotation + if _is_drop_fn(fn): + # Dropping potentially unsupported return type annotation for jit._drop + fn_def.returns = None # type:ignore[union-attr] + fn_def.type_comment = None # type:ignore[union-attr] + + # If MonkeyType is installed, get all the consolidated type traces + # for the arguments from type_trace_db + type_trace_db = torch.jit._script._get_type_trace_db() + pdt_arg_types = None + if monkeytype_trace and not isinstance(fn, _ParsedDef): # type: ignore[truthy-function] + qualname = get_qualified_name(fn) + pdt_arg_types = type_trace_db.get_args_types(qualname) + + return build_def( + parsed_def.ctx, + fn_def, + type_line, + def_name, + self_name=self_name, + pdt_arg_types=pdt_arg_types, + ) + + +# TODO: more robust handling of recognizing ignore context manager +def is_torch_jit_ignore_context_manager(stmt): + # checks if the statement is torch.jit.ignore context manager + if isinstance(stmt.items[0].context_expr, ast.Call): + # extract torch part + function = stmt.items[0].context_expr.func + if isinstance(function, ast.Attribute): + attr_name = function.attr + attr_value = function.value + if attr_name == "_IgnoreContextManager" and isinstance( + attr_value, ast.Attribute + ): + # there should be at most two nested attributes (e.g torch.jit._IgnoreContextManager) + if attr_value.attr == "jit" and isinstance(attr_value.value, ast.Name): + if attr_value.value.id == "torch": + return True + return False + + +class Builder: + def __call__(self, ctx, node): + method = getattr(self, "build_" + node.__class__.__name__, None) + if method is None: + raise UnsupportedNodeError(ctx, node) + return method(ctx, node) + + +def build_class_def(ctx, py_def, methods, properties, self_name, assigns): + r = ctx.make_range( + py_def.lineno, py_def.col_offset, py_def.col_offset + len("class") + ) + return ClassDef( + Ident(r, self_name), [Stmt(method) for method in methods], properties, assigns + ) + + +def build_def(ctx, py_def, type_line, def_name, self_name=None, pdt_arg_types=None): + body = py_def.body + r = ctx.make_range(py_def.lineno, py_def.col_offset, py_def.col_offset + len("def")) + + param_list = build_param_list(ctx, py_def.args, self_name, pdt_arg_types) + return_type = None + if getattr(py_def, "returns", None) is not None: + return_type = build_expr(ctx, py_def.returns) + + decl = Decl(r, param_list, return_type) + is_method = self_name is not None + if type_line is not None: + type_comment_decl = torch._C.parse_type_comment(type_line) + decl = torch._C.merge_type_from_type_comment(decl, type_comment_decl, is_method) + + return Def(Ident(r, def_name), decl, build_stmts(ctx, body)) + + +_vararg_kwarg_err = ( + "Compiled functions can't take variable number of arguments " + "or use keyword-only arguments with defaults" +) + + +def build_param_list(ctx, py_args, self_name, pdt_arg_types=None): + if py_args.kwarg is not None: + expr = py_args.kwarg + ctx_range = ctx.make_range( + expr.lineno, expr.col_offset - 1, expr.col_offset + len(expr.arg) + ) + raise NotSupportedError(ctx_range, _vararg_kwarg_err) + if py_args.vararg is not None: + expr = py_args.vararg + ctx_range = ctx.make_range( + expr.lineno, expr.col_offset - 1, expr.col_offset + len(expr.arg) + ) + raise NotSupportedError(ctx_range, _vararg_kwarg_err) + if len(py_args.kw_defaults) > 0: + # kw_defaults is a list of the values for the kwargs (which default to None), + # so they don't actually have line numbers. + for arg in py_args.kw_defaults: + if arg is not None: + ctx_range = build_expr(ctx, arg).range() + raise NotSupportedError(ctx_range, _vararg_kwarg_err) + + # List of Tuple of args and type as inferred by profile directed typing + arg_and_types = [ + ( + arg, + pdt_arg_types[arg.arg] + if pdt_arg_types and bool(pdt_arg_types[arg.arg]) + else None, + ) + for arg in py_args.args + ] + arg_and_types_kwonlyargs = [ + ( + arg, + pdt_arg_types[arg.arg] + if pdt_arg_types and bool(pdt_arg_types[arg.arg]) + else None, + ) + for arg in py_args.kwonlyargs + ] + + result = [ + build_param(ctx, arg, self_name, kwarg_only=False, pdt_arg_type=arg_type) + for arg, arg_type in arg_and_types + ] + result += [ + build_param(ctx, arg, self_name, kwarg_only=True, pdt_arg_type=arg_type) + for arg, arg_type in arg_and_types_kwonlyargs + ] + return result + + +def build_param(ctx, py_arg, self_name, kwarg_only, pdt_arg_type=None): + # NB: In Python3 py_arg is a pair of (str arg, expr? annotation) + name = py_arg.arg + r = ctx.make_range(py_arg.lineno, py_arg.col_offset, py_arg.col_offset + len(name)) + if getattr(py_arg, "annotation", None) is not None: + annotation_expr = build_expr(ctx, py_arg.annotation) + elif pdt_arg_type: + annotation_expr = Var(Ident(r, pdt_arg_type)) + elif self_name is not None and name == "self": + annotation_expr = Var(Ident(r, self_name)) + else: + annotation_expr = EmptyTypeAnnotation(r) + return Param(annotation_expr, Ident(r, name), kwarg_only) + + +def build_ignore_context_manager(ctx, stmt): + InputType = namedtuple("InputType", ["name", "ann"]) + OutputType = namedtuple("OutputType", ["name", "ann"]) + + def process_ins_outs(args): + # parse the context manager to figure out inputs and outputs + # with their annotated types + # TODO: add input, output validator + inputs = [] + outputs = [] + for arg in args: + var_name = arg.arg + var_ann = arg.value.value + var_decl_type, var_ann = var_ann.split(":") + if var_decl_type == "inp": + inputs.append(InputType(var_name, var_ann)) + if var_decl_type == "out": + outputs.append(OutputType(var_name, var_ann)) + return inputs, outputs + + def create_unique_name_ext(ctx, stmt): + # extension will be based on the full path filename plus + # the line number of original context manager + fn = re.sub(r"[^a-zA-Z0-9_]", "_", ctx.filename) + return f"{fn}_{stmt.lineno}" + + def build_return_ann_stmt(outputs): + return_type_ann = "" + return_statement_str = "return " + if len(outputs) == 0: + return_type_ann += " -> None" + if len(outputs) == 1: + return_type_ann = " -> " + outputs[0].ann + return_statement_str += outputs[0].name + if len(outputs) > 1: + return_type_ann = " -> Tuple" + return_type_ann += "[" + ", ".join([var.ann for var in outputs]) + "]" + return_statement_str += ", ".join([var.name for var in outputs]) + return return_type_ann, return_statement_str + + def build_args(args): + return ", ".join([arg.name for arg in args]) + + inputs, outputs = process_ins_outs(stmt.items[0].context_expr.keywords) + + # build the replacement function str with given inputs and outputs + ignore_function_name = "func_ignore_" + create_unique_name_ext(ctx, stmt) + ignore_function_str = "\ndef " + ignore_function_name + ignore_function_str += ( + "(" + ", ".join([var.name + " :" + var.ann for var in inputs]) + ")" + ) + + return_ann, return_stmt = build_return_ann_stmt(outputs) + ignore_function_str += return_ann + ": pass" + + # first create the functionDef object from just declaration + ignore_function = ast.parse(ignore_function_str).body[0] + + # dump the body of context manager to dummy function + ignore_function.body = stmt.body # type: ignore[attr-defined] + + # insert return statement to the function + return_stmt = ast.parse(return_stmt).body[0] + ignore_function.body.append(return_stmt) # type: ignore[attr-defined] + + # registers the custom function in the global context + ignore_func_str = "@torch.jit.ignore\n" + astunparse.unparse(ignore_function) + ignore_func_str += f'\nglobals()["{ignore_function_name}"] = {ignore_function_name}' + exec(ignore_func_str) # noqa: P204 + + # build the statements as: + # , , ... = torch.jit.frontend.(, ) + assign_str_lhs = build_args(outputs) + # this function will be registered in torch.jit.frontend module by default + assign_str_rhs = ( + f"torch.jit.frontend.{ignore_function_name}(" + build_args(inputs) + ")" + ) + + if len(outputs) > 0: + assign_str = assign_str_lhs + " = " + assign_str_rhs + else: + assign_str = assign_str_rhs + assign_ast = ast.parse(assign_str).body[0] + return assign_ast + + +def get_default_args(fn): + """ + Get a dictionary of default arguments for a function. + + Args: + fn: Callable - The function to inspect for default arguments. + Returns: + (Dict[str, Any]): mapping argument names to their default values if + :attr:`fn` is not None, else empty dictionary. + """ + if fn is None: + return {} + + signature = inspect.signature(fn) + + return { + k: v.default + for k, v in signature.parameters.items() + if v.default is not inspect.Parameter.empty + } + + +def get_default_args_for_class(cls): + """ + Get default arguments for all methods in a class (except for static methods). + + Args: + cls: type - The class type to inspect for default arguments. + Returns: + A Dict[str, Dict[str, Any]] which maps each method name to a Dict[str, Any] + that maps each argument name to its default value. + """ + # Get methods (except static methods because those are compiled separately as + # if they were independent script functions). + methods = inspect.getmembers( + cls, + predicate=lambda m: (inspect.ismethod(m) or inspect.isfunction(m)) + and not is_static_fn(cls, m.__name__) + and m.__name__ in cls.__dict__, + ) + + # Get method defaults. Property defaults do not need to be considered + # because setters cannot be invoked without a value. + defaults = { + method_name: get_default_args(method_impl) + for method_name, method_impl in methods + } + + return defaults + + +class WithItemBuilder(Builder): + @staticmethod + def build_withitem(ctx, item): + lineno = item.context_expr.lineno + start = item.context_expr.col_offset + end = start + len(pretty_node_names[ast.With]) + op_vars = item.optional_vars + r = ctx.make_range(lineno, start, end) + + return WithItem( + r, + build_expr(ctx, item.context_expr), + build_expr(ctx, op_vars) if op_vars else None, + ) + + +class StmtBuilder(Builder): + augassign_map = { + ast.Add: "+", + ast.Sub: "-", + ast.Mult: "*", + ast.Div: "/", + ast.Mod: "%", + ast.BitOr: "|", + ast.BitAnd: "&", + ast.BitXor: "^", + ast.LShift: "<<", + ast.RShift: ">>", + ast.Pow: "**", + } + + @staticmethod + def build_Expr(ctx, stmt): + value = stmt.value + if value.__class__.__name__ == "Str": + # If a statement is a string literal expression, + # then it is a docstring. Just ignore it. + return None + else: + return ExprStmt(build_expr(ctx, value)) + + @staticmethod + def build_Assign(ctx, stmt): + rhs = build_expr(ctx, stmt.value) + lhs = [build_expr(ctx, x) for x in stmt.targets] + return Assign(lhs, rhs) + + @staticmethod + def build_AnnAssign(ctx, stmt): + if stmt.value is None: + raise UnsupportedNodeError(ctx, stmt, reason="without assigned value") + + # Disallow type annotations on instance attributes outside of __init__ + if ( + type(stmt.target) == ast.Attribute + and stmt.target.value.id == "self" # type: ignore[attr-defined] + and ctx.funcname != "__init__" + ): + start = stmt.col_offset + end = start + len(f"self.{stmt.target.attr}") + if hasattr(stmt.annotation, "id"): + end += len(f": {stmt.annotation.id}") + sr = ctx.make_range(stmt.lineno, start, end) + raise ValueError( + "Type annotations on instance attributes must be declared in " + f"__init__, not '{ctx.funcname}': {sr}" + ) + + rhs = build_expr(ctx, stmt.value) + lhs = build_expr(ctx, stmt.target) + the_type = build_expr(ctx, stmt.annotation) + return Assign([lhs], rhs, the_type) + + @staticmethod + def build_Delete(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("del")) + + return Delete(r, [build_expr(ctx, target) for target in stmt.targets]) + + @staticmethod + def build_Return(ctx, stmt): + r = ctx.make_range( + stmt.lineno, stmt.col_offset, stmt.col_offset + len("return") + ) + return Return(r, None if stmt.value is None else build_expr(ctx, stmt.value)) + + @staticmethod + def build_Raise(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("raise")) + expr = build_expr(ctx, stmt.exc) + return Raise(r, expr) + + @staticmethod + def build_Assert(ctx, stmt): + r = ctx.make_range( + stmt.lineno, stmt.col_offset, stmt.col_offset + len("assert") + ) + test = build_expr(ctx, stmt.test) + msg = build_expr(ctx, stmt.msg) if stmt.msg is not None else None + return Assert(r, test, msg) + + @staticmethod + def build_AugAssign(ctx, stmt): + lhs = build_expr(ctx, stmt.target) + rhs = build_expr(ctx, stmt.value) + op = type(stmt.op) + if op in StmtBuilder.augassign_map: + op_token = StmtBuilder.augassign_map[op] + else: + raise NotSupportedError( + find_before(ctx, rhs.range().start, "=", offsets=(-1, 0)), + "unsupported kind of augmented assignment: " + op.__name__, + ) + return AugAssign(lhs, op_token, rhs) + + @staticmethod + def build_While(ctx, stmt): + if stmt.orelse: + # TODO: try to recover the location of else:? Python doesn't give us useful + # annotations in this case + raise NotSupportedError( + None, "else branches of while loops aren't supported" + ) + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("while")) + return While(r, build_expr(ctx, stmt.test), build_stmts(ctx, stmt.body)) + + @staticmethod + def build_For(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("for")) + if stmt.orelse: + raise NotSupportedError(r, "else branches of for loops aren't supported") + + return For( + r, + [build_expr(ctx, stmt.target)], + [build_expr(ctx, stmt.iter)], + build_stmts(ctx, stmt.body), + ) + + @staticmethod + def build_If(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("if")) + return If( + r, + build_expr(ctx, stmt.test), + build_stmts(ctx, stmt.body), + build_stmts(ctx, stmt.orelse), + ) + + @staticmethod + def build_Print(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("print")) + if stmt.dest: + raise NotSupportedError( + r, "print statements with non-default destinations aren't supported" + ) + args = [build_expr(ctx, val) for val in stmt.values] + return ExprStmt(Apply(Var(Ident(r, "print")), args, [])) + + @staticmethod + def build_Pass(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("pass")) + return Pass(r) + + @staticmethod + def build_Break(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("break")) + return Break(r) + + @staticmethod + def build_Continue(ctx, stmt): + r = ctx.make_range( + stmt.lineno, stmt.col_offset, stmt.col_offset + len("continue") + ) + return Continue(r) + + @staticmethod + def build_With(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("with")) + # Handle ignore context manager + if is_torch_jit_ignore_context_manager(stmt): + if not _IS_ASTUNPARSE_INSTALLED: + raise RuntimeError( + "torch.jit._IgnoreContextManager requires installing Python library `astunparse`, \ + please install it in your Python environment" + ) + assign_ast = build_ignore_context_manager(ctx, stmt) + return build_stmt(ctx, assign_ast) + return With(r, build_withitems(ctx, stmt.items), build_stmts(ctx, stmt.body)) + + +class ExprBuilder(Builder): + binop_map = { + ast.Add: "+", + ast.Sub: "-", + ast.Mult: "*", + ast.Div: "/", + ast.Pow: "**", + ast.Mod: "%", + ast.FloorDiv: "//", + ast.BitAnd: "&", + ast.BitXor: "^", + ast.BitOr: "|", + ast.LShift: "<<", + ast.RShift: ">>", + } + + binop_map[ast.MatMult] = "@" + + unop_map = { + ast.Not: "not", + ast.USub: "-", + ast.Invert: "~", + } + + boolop_map = { + ast.And: "and", + ast.Or: "or", + } + + cmpop_map = { + ast.Eq: "==", + ast.NotEq: "!=", + ast.LtE: "<=", + ast.Lt: "<", + ast.GtE: ">=", + ast.Gt: ">", + ast.Is: "is", + ast.IsNot: "is not", + ast.In: "in", + ast.NotIn: "not in", + } + + @staticmethod + def build_Attribute(ctx, expr): + base = build_expr(ctx, expr.value) + # expr.attr is just a string, so it's not annotated in any way, so we have + # to build the range manually + source = ctx.source.encode("utf-8") + + def get_char(index): + return chr(source[index]) + + start_pos = base.range().end + 1 + while get_char(start_pos) in string.whitespace: # Skip whitespace + start_pos += 1 + end_pos = start_pos + len(expr.attr) + name_range = ctx.make_raw_range(start_pos, end_pos) + return Select(base, Ident(name_range, expr.attr)) + + @staticmethod + def build_Call(ctx, expr): + func = build_expr(ctx, expr.func) + args = [build_expr(ctx, py_arg) for py_arg in expr.args] + if hasattr(expr, "starargs") and expr.starargs: + stararg_expr = build_expr(ctx, expr.starargs) + args += [Starred(stararg_expr.range(), stararg_expr)] + kwargs = [] + for kw in expr.keywords: + kw_expr = build_expr(ctx, kw.value) + # XXX: we could do a better job at figuring out the range for the name here + if not kw.arg: + raise NotSupportedError( + kw_expr.range(), "keyword-arg expansion is not supported" + ) + kwargs.append(Attribute(Ident(kw_expr.range(), kw.arg), kw_expr)) + return Apply(func, args, kwargs) + + @staticmethod + def build_Ellipsis(ctx, expr): + r = ctx.make_range( + expr.lineno, expr.col_offset, expr.col_offset + 3 + ) # len("...") == 3 + return Dots(r) + + @staticmethod + def build_Name(ctx, expr): + r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(expr.id)) + if expr.id.startswith(_reserved_prefix): + raise NotSupportedError( + r, + "names of variables used in JIT-ed functions " + "can't start with " + _reserved_prefix, + ) + if expr.id == "True": + return TrueLiteral(r) + elif expr.id == "False": + return FalseLiteral(r) + elif expr.id == "None": + return NoneLiteral(r) + elif expr.id == "Ellipsis": + return Dots(r) + return Var(Ident(r, expr.id)) + + @staticmethod + def build_NameConstant(ctx, expr): + r = ctx.make_range( + expr.lineno, expr.col_offset, expr.col_offset + len(str(expr.value)) + ) + if expr.value is True: + return TrueLiteral(r) + elif expr.value is False: + return FalseLiteral(r) + elif expr.value is None: + return NoneLiteral(r) + elif expr.value == Ellipsis: + return Dots(r) + else: + raise ValueError("Name constant value unsupported: " + str(expr.value)) + + @staticmethod + def build_BinOp(ctx, expr): + lhs = build_expr(ctx, expr.left) + rhs = build_expr(ctx, expr.right) + op = type(expr.op) + + if op == ast.Div and not ctx.uses_true_division: + err_range = ctx.make_raw_range(lhs.range().end, rhs.range().start) + raise FrontendError( + err_range, + "Division of ints in TorchScript uses Python 3 true " + "division semantics. Please put `from __future__ " + "import division` at the top of your file", + ) + op_token = ExprBuilder.binop_map.get(op) + if op_token is None: + err_range = ctx.make_raw_range(lhs.range().end, rhs.range().start) + raise NotSupportedError( + err_range, "unsupported binary operator: " + op.__name__ + ) + return BinOp(op_token, lhs, rhs) + + @staticmethod + def build_UnaryOp(ctx, expr): + sub_expr = build_expr(ctx, expr.operand) + op = type(expr.op) + op_token = ExprBuilder.unop_map.get(op) + if op_token is None: + raise NotSupportedError( + expr.range(), "unsupported unary operator: " + op.__name__ + ) + r = ctx.make_range( + expr.lineno, expr.col_offset, expr.col_offset + len(op_token) + ) + return UnaryOp(r, op_token, sub_expr) + + @staticmethod + def build_BoolOp(ctx, expr): + if len(expr.values) < 2: + raise AssertionError( + "expected at least 2 values in BoolOp, but got " + str(len(expr.values)) + ) + sub_exprs = [build_expr(ctx, sub_expr) for sub_expr in expr.values] + op = type(expr.op) + op_token = ExprBuilder.boolop_map.get(op) + if op_token is None: + err_range = ctx.make_raw_range( + sub_exprs[0].range().end, sub_exprs[1].range().start + ) + raise NotSupportedError( + err_range, "unsupported boolean operator: " + op.__name__ + ) + lhs = sub_exprs[0] + for rhs in sub_exprs[1:]: + lhs = BinOp(op_token, lhs, rhs) + return lhs + + @staticmethod + def build_IfExp(ctx, expr): + return TernaryIf( + build_expr(ctx, expr.test), + build_expr(ctx, expr.body), + build_expr(ctx, expr.orelse), + ) + + @staticmethod + def build_Compare(ctx, expr): + operands = [build_expr(ctx, e) for e in [expr.left] + list(expr.comparators)] + result = None + for lhs, op_, rhs in zip(operands, expr.ops, operands[1:]): + op = type(op_) + op_token = ExprBuilder.cmpop_map.get(op) + r = ctx.make_raw_range(lhs.range().end, rhs.range().start) + if op_token is None: + raise NotSupportedError( + r, "unsupported comparison operator: " + op.__name__ + ) + + if op == ast.NotIn: + # NB: `not in` is just `not( in )`, so we don't introduce new tree view + # but just make it a nested call in our tree view structure + in_expr = BinOp("in", lhs, rhs) + cmp_expr = UnaryOp(r, "not", in_expr) + else: + cmp_expr = BinOp(op_token, lhs, rhs) + + if result is None: + result = cmp_expr + else: + result = BinOp("and", result, cmp_expr) + return result + + @staticmethod + def build_Subscript(ctx, expr): + def build_SliceExpr(ctx, base, slice_expr): + lower = ( + build_expr(ctx, slice_expr.lower) + if slice_expr.lower is not None + else None + ) + upper = ( + build_expr(ctx, slice_expr.upper) + if slice_expr.upper is not None + else None + ) + step = ( + build_expr(ctx, slice_expr.step) + if slice_expr.step is not None + else None + ) + return SliceExpr(base.range(), lower, upper, step) + + def build_Index(ctx, base, index_expr): + if isinstance(index_expr.value, ast.Tuple): + raise NotSupportedError( + base.range(), + "slicing multiple dimensions with tuples not supported yet", + ) + return build_expr(ctx, index_expr.value) + + def build_ExtSlice(ctx, base, extslice): + sub_exprs = [] + for expr in extslice.dims: + sub_type = type(expr) + if sub_type is ast.Index: + sub_exprs.append(build_Index(ctx, base, expr)) + elif sub_type is ast.Slice: + sub_exprs.append(build_SliceExpr(ctx, base, expr)) + elif sub_type is ast.Constant and expr.value is Ellipsis: + sub_exprs.append(Dots(base.range())) + else: + raise NotSupportedError( + base.range(), + f"slicing multiple dimensions with {sub_type} not supported", + ) + return sub_exprs + + base = build_expr(ctx, expr.value) + sub_type = type(expr.slice) + if sub_type is ast.Index: + if isinstance(expr.slice.value, ast.Tuple): + # N-dimensional indexing using Tuple: x[(i, j, k)] is equivalent to x[i, j, k] + # XXX: Indexing using a list is **different**! It triggers advanced indexing. + indices = [ + build_expr(ctx, index_expr) for index_expr in expr.slice.value.elts + ] + if not indices: + # `col_offset` is an int, but `end_col_offset` is + # `Optional[int]`. The magic number is here to make + # sure we can parse `()` on any machine + r = ctx.make_range( + expr.lineno, + expr.slice.value.col_offset, + expr.slice.value.col_offset + 2, + ) + tup = TupleLiteral(r, []) + indices.append(tup) + return Subscript(base, indices) + else: + return Subscript(base, [build_expr(ctx, expr.slice.value)]) + elif sub_type is ast.Slice: + return Subscript(base, [build_SliceExpr(ctx, base, expr.slice)]) + elif sub_type is ast.ExtSlice: + return Subscript(base, build_ExtSlice(ctx, base, expr.slice)) + elif sys.version_info >= ( + 3, + 9, + ): # In Python3.9 array indicies are not wrapped in ast.Index + if sub_type is ast.Tuple: + # N-dimensional indexing using Tuple: x[(i, j, k)] is equivalent to x[i, j, k] + indices = [] + for index_expr in expr.slice.elts: + if isinstance(index_expr, ast.Slice): + indices.append(build_SliceExpr(ctx, base, index_expr)) + else: + indices.append(build_expr(ctx, index_expr)) + # Special-case logic for `typing.Tuple[()]` + if not indices: + # See note above r.e. magic number + r = ctx.make_range( + expr.lineno, expr.slice.col_offset, expr.slice.col_offset + 2 + ) + tup = TupleLiteral(r, []) + indices.append(tup) + return Subscript(base, indices) + return Subscript(base, [build_expr(ctx, expr.slice)]) + else: # Ellipsis (can only happen in Python 2) + raise NotSupportedError(base.range(), "ellipsis is not supported") + + @staticmethod + def build_List(ctx, expr): + return ListLiteral( + ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1), + [build_expr(ctx, e) for e in expr.elts], + ) + + @staticmethod + def build_Tuple(ctx, expr): + return TupleLiteral( + ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1), + [build_expr(ctx, e) for e in expr.elts], + ) + + @staticmethod + def build_Dict(ctx, expr): + range = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1) + if expr.keys and not expr.keys[0]: + raise NotSupportedError( + range, "Dict expansion (e.g. `{**dict}`) is not supported" + ) + return DictLiteral( + range, + [build_expr(ctx, e) for e in expr.keys], + [build_expr(ctx, e) for e in expr.values], + ) + + @staticmethod + def build_Num(ctx, expr): + value = str(expr.value) + r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(value)) + return Const(r, value) + + @staticmethod + def build_Constant(ctx, expr): + value = expr.value + if value is None or isinstance(value, bool): + # NB: this check has to happen before the int check because bool is + # a subclass of int + return ExprBuilder.build_NameConstant(ctx, expr) + if isinstance(value, (int, float, complex)): + return ExprBuilder.build_Num(ctx, expr) + elif isinstance(value, str): + return ExprBuilder.build_Str(ctx, expr) + elif isinstance(value, type(Ellipsis)): + return ExprBuilder.build_Ellipsis(ctx, expr) + else: + error_range = ctx.make_range( + expr.lineno, expr.col_offset, expr.col_offset + len(str(value)) + ) + raise FrontendError(error_range, "Unknown Constant expression type") + + @staticmethod + def build_Str(ctx, expr): + value = str(expr.value) + r = ctx.make_range( + expr.lineno, expr.col_offset, expr.col_offset + len(value) + 1 + ) + return StringLiteral(r, value) + + @staticmethod + def build_JoinedStr(ctx, expr): + s = "" + args = [] + for value in expr.values: + r = ctx.make_range(value.lineno, value.col_offset, value.col_offset + 1) + if isinstance(value, ast.FormattedValue): + if value.conversion != -1: + raise NotSupportedError(r, "Don't support conversion in JoinedStr") + if value.format_spec is not None: + raise NotSupportedError(r, "Don't support formatting in JoinedStr") + s += "{}" + args.append(build_expr(ctx, value.value)) + elif isinstance(value, ast.Constant): + s += value.value + else: + raise NotSupportedError(r, "Unsupported value in JoinedStr") + + r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1) + return Apply(Select(StringLiteral(r, s), Ident(r, "format")), args, []) + + @staticmethod + def build_ListComp(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset) + if len(stmt.generators) != 1: + raise NotSupportedError(r, "Only a single generator is currently supported") + + if len(stmt.generators[0].ifs) != 0: + raise NotSupportedError(r, "Comprehension ifs are not supported yet") + + elt_expr = build_expr(ctx, stmt.elt) + target_expr = build_expr(ctx, stmt.generators[0].target) + iter_expr = build_expr(ctx, stmt.generators[0].iter) + + return ListComp(r, elt_expr, target_expr, iter_expr) + + @staticmethod + def build_GeneratorExp(ctx, stmt): + # Convert Generator expression to ListComp + return ExprBuilder.build_ListComp(ctx, stmt) + + @staticmethod + def build_DictComp(ctx, stmt): + r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset) + if len(stmt.generators) != 1: + raise NotSupportedError(r, "Only a single generator is currently supported") + + if len(stmt.generators[0].ifs) != 0: + raise NotSupportedError(r, "Comprehension ifs are not supported yet") + + key_expr = build_expr(ctx, stmt.key) + value_expr = build_expr(ctx, stmt.value) + target_expr = build_expr(ctx, stmt.generators[0].target) + iter_expr = build_expr(ctx, stmt.generators[0].iter) + + return DictComp(r, key_expr, value_expr, target_expr, iter_expr) + + @staticmethod + def build_Starred(ctx, expr): + r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1) + return Starred(r, build_expr(ctx, expr.value)) + + +build_expr = ExprBuilder() +build_stmt = StmtBuilder() +build_withitem = WithItemBuilder() + + +def find_before(ctx, pos, substr, offsets=(0, 0)): + new_pos = ctx.source[:pos].rindex(substr) + return ctx.make_raw_range(new_pos + offsets[0], new_pos + len(substr) + offsets[1]) diff --git a/lib/python3.12/site-packages/torch/jit/generate_bytecode.py b/lib/python3.12/site-packages/torch/jit/generate_bytecode.py new file mode 100644 index 0000000000000000000000000000000000000000..f66bf7bfc4c1b25a35b342ab6792af025266a3a8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/generate_bytecode.py @@ -0,0 +1,34 @@ +# mypy: allow-untyped-defs +from typing import List + +from torch._C import _compile_graph_to_code_table, _generate_upgraders_graph + + +def format_bytecode(table): + # given a nested tuple, convert it to nested list + def listify(content): + if not isinstance(content, tuple): + return content + return [listify(i) for i in content] + + formatted_table = {} + for entry in table: + identifier = entry[0] + content = entry[1] + content = listify(content) + formatted_table[identifier] = content + return formatted_table + + +def generate_upgraders_bytecode() -> List: + yaml_content = [] + upgraders_graph_map = _generate_upgraders_graph() + for upgrader_name, upgrader_graph in upgraders_graph_map.items(): + bytecode_table = _compile_graph_to_code_table(upgrader_name, upgrader_graph) + entry = {upgrader_name: format_bytecode(bytecode_table)} + yaml_content.append(entry) + return yaml_content + + +if __name__ == "__main__": + raise RuntimeError("This file is not meant to be run directly") diff --git a/lib/python3.12/site-packages/torch/jit/mobile/__init__.py b/lib/python3.12/site-packages/torch/jit/mobile/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..593dabf2fc43ca075d604977a0d254ea02f3ba55 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/mobile/__init__.py @@ -0,0 +1,232 @@ +# mypy: allow-untyped-defs +import os + +import torch +from torch.jit._serialization import validate_map_location + + +def _load_for_lite_interpreter(f, map_location=None): + r""" + Load a :class:`LiteScriptModule` saved with :func:`torch.jit._save_for_lite_interpreter`. + + Args: + f: a file-like object (has to implement read, readline, tell, and seek), + or a string containing a file name + map_location: a string or torch.device used to dynamically remap + storages to an alternative set of devices. + + Returns: + A :class:`LiteScriptModule` object. + + Example: + + .. testcode:: + + import torch + import io + + # Load LiteScriptModule from saved file path + torch.jit._load_for_lite_interpreter('lite_script_module.pt') + + # Load LiteScriptModule from io.BytesIO object + with open('lite_script_module.pt', 'rb') as f: + buffer = io.BytesIO(f.read()) + + # Load all tensors to the original device + torch.jit.mobile._load_for_lite_interpreter(buffer) + """ + if isinstance(f, (str, os.PathLike)): + if not os.path.exists(f): + raise ValueError(f"The provided filename {f} does not exist") + if os.path.isdir(f): + raise ValueError(f"The provided filename {f} is a directory") + + map_location = validate_map_location(map_location) + + if isinstance(f, (str, os.PathLike)): + cpp_module = torch._C._load_for_lite_interpreter(os.fspath(f), map_location) + else: + cpp_module = torch._C._load_for_lite_interpreter_from_buffer( + f.read(), map_location + ) + + return LiteScriptModule(cpp_module) + + +class LiteScriptModule: + def __init__(self, cpp_module): + self._c = cpp_module + super().__init__() + + def __call__(self, *input): + return self._c.forward(input) + + def find_method(self, method_name): + return self._c.find_method(method_name) + + def forward(self, *input): + return self._c.forward(input) + + def run_method(self, method_name, *input): + return self._c.run_method(method_name, input) + + +def _export_operator_list(module: LiteScriptModule): + r"""Return a set of root operator names (with overload name) that are used by any method in this mobile module.""" + return torch._C._export_operator_list(module._c) + + +def _get_model_bytecode_version(f_input) -> int: + r"""Take a file-like object to return an integer. + + Args: + f_input: a file-like object (has to implement read, readline, tell, and seek), + or a string containing a file name + + Returns: + version: An integer. If the integer is -1, the version is invalid. A warning + will show in the log. + + Example: + .. testcode:: + + from torch.jit.mobile import _get_model_bytecode_version + + # Get bytecode version from a saved file path + version = _get_model_bytecode_version("path/to/model.ptl") + + """ + if isinstance(f_input, (str, os.PathLike)): + if not os.path.exists(f_input): + raise ValueError(f"The provided filename {f_input} does not exist") + if os.path.isdir(f_input): + raise ValueError(f"The provided filename {f_input} is a directory") + + if isinstance(f_input, (str, os.PathLike)): + return torch._C._get_model_bytecode_version(os.fspath(f_input)) + else: + return torch._C._get_model_bytecode_version_from_buffer(f_input.read()) + + +def _get_mobile_model_contained_types(f_input) -> int: + r"""Take a file-like object and return a set of string, like ("int", "Optional"). + + Args: + f_input: a file-like object (has to implement read, readline, tell, and seek), + or a string containing a file name + + Returns: + type_list: A set of string, like ("int", "Optional"). These are types used in bytecode. + + Example: + + .. testcode:: + + from torch.jit.mobile import _get_mobile_model_contained_types + + # Get type list from a saved file path + type_list = _get_mobile_model_contained_types("path/to/model.ptl") + + """ + if isinstance(f_input, (str, os.PathLike)): + if not os.path.exists(f_input): + raise ValueError(f"The provided filename {f_input} does not exist") + if os.path.isdir(f_input): + raise ValueError(f"The provided filename {f_input} is a directory") + + if isinstance(f_input, (str, os.PathLike)): + return torch._C._get_mobile_model_contained_types(os.fspath(f_input)) + else: + return torch._C._get_mobile_model_contained_types_from_buffer(f_input.read()) + + +def _backport_for_mobile(f_input, f_output, to_version): + r"""Take a input string containing a file name (file-like object) and a new destination to return a boolean. + + Args: + f_input: a file-like object (has to implement read, readline, tell, and seek), + or a string containing a file name + f_output: path to new model destination + to_version: the expected output model bytecode version + Returns: + success: A boolean. If backport success, return true, otherwise false + """ + if isinstance(f_input, (str, os.PathLike)): + if not os.path.exists(f_input): + raise ValueError(f"The provided filename {f_input} does not exist") + if os.path.isdir(f_input): + raise ValueError(f"The provided filename {f_input} is a directory") + + if (isinstance(f_input, (str, os.PathLike))) and ( + isinstance(f_output, (str, os.PathLike)) + ): + return torch._C._backport_for_mobile( + os.fspath(f_input), os.fspath(f_output), to_version + ) + else: + return torch._C._backport_for_mobile_from_buffer( + f_input.read(), str(f_output), to_version + ) + + +def _backport_for_mobile_to_buffer(f_input, to_version): + r"""Take a string containing a file name (file-like object). + + Args: + f_input: a file-like object (has to implement read, readline, tell, and seek), + or a string containing a file name + + """ + if isinstance(f_input, (str, os.PathLike)): + if not os.path.exists(f_input): + raise ValueError(f"The provided filename {f_input} does not exist") + if os.path.isdir(f_input): + raise ValueError(f"The provided filename {f_input} is a directory") + + if isinstance(f_input, (str, os.PathLike)): + return torch._C._backport_for_mobile_to_buffer(os.fspath(f_input), to_version) + else: + return torch._C._backport_for_mobile_from_buffer_to_buffer( + f_input.read(), to_version + ) + + +def _get_model_ops_and_info(f_input): + r"""Retrieve the root (top level) operators of a model and their corresponding compatibility info. + + These root operators can call other operators within them (traced ops), and + a root op can call many different traced ops depending on internal code paths in the root op. + These traced ops are not returned by this function. Those operators are abstracted into the + runtime as an implementation detail (and the traced ops themselves can also call other operators) + making retrieving them difficult and their value from this api negligible since they will differ + between which runtime version the model is run on. Because of this, there is a false positive this + api can't prevent in a compatibility usecase. All the root ops of a model are present in a + target runtime, but not all the traced ops are which prevents a model from being able to run. + Args: + f_input: a file-like object (has to implement read, readline, tell, and seek), + or a string containing a file name + + Returns: + Operators and info: A Dictionary mapping strings (the qualified names of the root operators) + of the model to their OperatorInfo structs. + + Example: + + .. testcode:: + + from torch.jit.mobile import _get_model_ops_and_info + + # Get bytecode version from a saved file path + ops_and_info = _get_model_ops_and_info("path/to/model.ptl") + + """ + if isinstance(f_input, (str, os.PathLike)): + if not os.path.exists(f_input): + raise ValueError(f"The provided filename {f_input} does not exist") + if os.path.isdir(f_input): + raise ValueError(f"The provided filename {f_input} is a directory") + + if isinstance(f_input, (str, os.PathLike)): + return torch._C._get_model_ops_and_info(os.fspath(f_input)) + else: + return torch._C._get_model_ops_and_info(f_input.read()) diff --git a/lib/python3.12/site-packages/torch/jit/mobile/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/jit/mobile/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87f7934014813dee989bf7a23167199939f891b1 Binary files /dev/null and b/lib/python3.12/site-packages/torch/jit/mobile/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/jit/unsupported_tensor_ops.py b/lib/python3.12/site-packages/torch/jit/unsupported_tensor_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..903c8aafba26b22c04b8a6734e308ababb4c3c37 --- /dev/null +++ b/lib/python3.12/site-packages/torch/jit/unsupported_tensor_ops.py @@ -0,0 +1,78 @@ +# mypy: allow-untyped-defs +from textwrap import dedent +from typing import Any, Dict + +import torch.jit + + +def execWrapper(code, glob, loc): + exec(code, glob, loc) + + +def _gen_unsupported_methods_properties(): + tensor_attrs = set(filter(lambda x: x[0] != "_", dir(torch.Tensor))) + tensor = torch.tensor([2]) + funcs_template = dedent( + """ + def func(x): + return x.{op}() + """ + ) + + deprecated_apis = { + "volatile", + "resize", + "reinforce", + "new", + "name", + "map2_", + "has_names", + "grad_fn", + "resize_as", + } + tensor_attrs = tensor_attrs - deprecated_apis + + properties = [] + methods = [] + sorted_tensor_attrs = sorted(tensor_attrs, key=lambda x: x.lower()) + for attr in sorted_tensor_attrs: + funcs_str = funcs_template.format(op=attr) + scope: Dict[str, Any] = {} + execWrapper(funcs_str, globals(), scope) + try: + torch.jit.CompilationUnit(funcs_str) + except Exception as e: + if "nonexistent attribute" not in repr(e): + continue + attr_repr = repr(getattr(tensor, attr)) + if "bound method" in attr_repr or "built-in method" in attr_repr: + methods.append(attr) + else: + properties.append(attr) + + mapped_methods = ("\t* :meth:`~torch.Tensor." + x + r"`" for x in methods) + mapped_properties = ("\t* :attr:`~torch.Tensor." + x + r"`" for x in properties) + return "\n".join(mapped_methods), "\n".join(mapped_properties) + + +def _list_unsupported_tensor_ops(): + header = """\n\n +Unsupported Tensor Methods +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + """ + methods, properties = _gen_unsupported_methods_properties() + return ( + header + + "\n" + + methods + + """ + +Unsupported Tensor Properties +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + """ + + "\n" + + properties + ) + + +__doc__ = _list_unsupported_tensor_ops() diff --git a/lib/python3.12/site-packages/torch/lib/libcaffe2_nvrtc.so b/lib/python3.12/site-packages/torch/lib/libcaffe2_nvrtc.so new file mode 100644 index 0000000000000000000000000000000000000000..9a5baba80cda9f1785bb122830af65f95dd2be52 Binary files /dev/null and b/lib/python3.12/site-packages/torch/lib/libcaffe2_nvrtc.so differ diff --git a/lib/python3.12/site-packages/torch/lib/libshm.so b/lib/python3.12/site-packages/torch/lib/libshm.so new file mode 100644 index 0000000000000000000000000000000000000000..02efdf0b98d1d527c66b7e6b4e458d96847f1250 Binary files /dev/null and b/lib/python3.12/site-packages/torch/lib/libshm.so differ diff --git a/lib/python3.12/site-packages/torch/lib/libtorch_global_deps.so b/lib/python3.12/site-packages/torch/lib/libtorch_global_deps.so new file mode 100644 index 0000000000000000000000000000000000000000..e4e13e2cfb479c93129c488f67908fa0dafc68c0 Binary files /dev/null and b/lib/python3.12/site-packages/torch/lib/libtorch_global_deps.so differ diff --git a/lib/python3.12/site-packages/torch/masked/__init__.py b/lib/python3.12/site-packages/torch/masked/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d00ba1e8d5aff6a18490ac7b16a629ac36e3dcb5 --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/__init__.py @@ -0,0 +1,57 @@ +from torch.masked._ops import ( + _canonical_dim, + _combine_input_and_mask, + _generate_docstring, + _input_mask, + _output_mask, + _reduction_identity, + _where, + amax, + amin, + argmax, + argmin, + cumprod, + cumsum, + log_softmax, + logaddexp, + logsumexp, + mean, + median, + norm, + normalize, + prod, + softmax, + softmin, + std, + sum, + var, +) +from torch.masked.maskedtensor.core import is_masked_tensor, MaskedTensor +from torch.masked.maskedtensor.creation import as_masked_tensor, masked_tensor + + +__all__ = [ + "amax", + "amin", + "argmax", + "argmin", + "as_masked_tensor", + "cumprod", + "cumsum", + "is_masked_tensor", + "log_softmax", + "logaddexp", + "logsumexp", + "masked_tensor", + "MaskedTensor", + "mean", + "median", + "norm", + "normalize", + "prod", + "softmax", + "softmin", + "std", + "sum", + "var", +] diff --git a/lib/python3.12/site-packages/torch/masked/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8cb4c0103c0a4b5110a3f6b42ae8382590c1a56 Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/__pycache__/_docs.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/__pycache__/_docs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05ba3d1bdbfb5444efd2cf462f9eda01698c0e8b Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/__pycache__/_docs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/__pycache__/_ops.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/__pycache__/_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..036c4fdbcb5d9f61c8714f431682d1db99ff41bb Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/__pycache__/_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/_docs.py b/lib/python3.12/site-packages/torch/masked/_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..fa130bbefbc5caa7373459ef2fc3dc5292239948 --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/_docs.py @@ -0,0 +1,1177 @@ +# This file is generated, do not modify it! +# +# To update this file, run the update masked docs script as follows: +# +# python tools/update_masked_docs.py +# +# The script must be called from an environment where the development +# version of torch package can be imported and is functional. +# + +amax_docstring = """amax(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor + +Returns maximum of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. + +The identity value of maximum operation, which is used to start the +reduction, depends on input dtype. For instance, for float32, uint8, +and int32 dtypes, the identity values are ``-inf``, ``0``, and ``-2147483648``, respectively. + +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in maximum computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of maximum operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.amax(input, 1, mask=mask) + tensor([ -1, -9223372036854775808]) +""" + +amin_docstring = """amin(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor + +Returns minimum of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. + +The identity value of minimum operation, which is used to start the +reduction, depends on input dtype. For instance, for float32, uint8, +and int32 dtypes, the identity values are ``inf``, ``255``, and ``2147483647``, respectively. + +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in minimum computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of minimum operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.amin(input, 1, mask=mask) + tensor([ -3, 9223372036854775807]) +""" + +argmax_docstring = """argmax(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor +Returns argmax of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. +The identity value of argmax operation, which is used to start the +reduction, depends on input dtype. For instance, for float32, uint8, +and int32 dtypes, the identity values are ``-inf``, ``0``, and ``-2147483648``, respectively. +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in argmax computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of argmax operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int): the dimension along which argmax is computed. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.argmax(input, 1, mask=mask) + tensor([2, 0]) +""" + +argmin_docstring = """argmin(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor +Returns argmin of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. +The identity value of argmin operation, which is used to start the +reduction, depends on input dtype. For instance, for float32, uint8, +and int32 dtypes, the identity values are ``inf``, ``255``, and ``2147483647``, respectively. +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in argmin computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of argmin operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int): the dimension along which argmin is computed. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.argmin(input, 1, mask=mask) + tensor([0, 0]) +""" + +cumprod_docstring = """cumprod(input, dim, *, dtype=None, mask=None) -> Tensor + +Returns cumulative_prod of all the slices in the :attr:`input` tensor +along :attr:`dim` while the :attr:`input` elements are masked out +according to the boolean tensor :attr:`mask`. + +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Cumsum of i-th element in ``x`` is +defined as ``prod(x[:i])``. + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True then +the corresponding element in :attr:`input` tensor will be included in +cumulative_prod computation, otherwise the element is ignored. + +The values of masked-out elements of the output tensor have undefined +value: it may or may not be set to zero or nan; the choice may correspond to +the value that leads to the most efficient storage of :attr:`output` +tensor. + +The mask of the cumulative_prod output tensor can be computed as +``torch.broadcast_to(mask, input.shape)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int): the dimension along which cumulative_prod is computed. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]]) + >>> input + tensor([[-3., -2., -1.], + [ 0., 1., 2.]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.cumprod(input, 1, mask=mask) + tensor([[-3., -3., 3.], + [ 1., 1., 1.]]) +""" + +cumsum_docstring = """cumsum(input, dim, *, dtype=None, mask=None) -> Tensor + +Returns cumulative_sum of all the slices in the :attr:`input` tensor +along :attr:`dim` while the :attr:`input` elements are masked out +according to the boolean tensor :attr:`mask`. + +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Cumsum of i-th element in ``x`` is +defined as ``sum(x[:i])``. + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True then +the corresponding element in :attr:`input` tensor will be included in +cumulative_sum computation, otherwise the element is ignored. + +The values of masked-out elements of the output tensor have undefined +value: it may or may not be set to zero or nan; the choice may correspond to +the value that leads to the most efficient storage of :attr:`output` +tensor. + +The mask of the cumulative_sum output tensor can be computed as +``torch.broadcast_to(mask, input.shape)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int): the dimension along which cumulative_sum is computed. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]]) + >>> input + tensor([[-3., -2., -1.], + [ 0., 1., 2.]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.cumsum(input, 1, mask=mask) + tensor([[-3., -3., -4.], + [ 0., 0., 0.]]) +""" + +log_softmax_docstring = """log_softmax(input, dim, *, dtype=None, mask=None) -> Tensor + +Returns log_softmax of all the slices in the :attr:`input` tensor +along :attr:`dim` while the :attr:`input` elements are masked out +according to the boolean tensor :attr:`mask`. + +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. LogSoftmax of i-th element in ``x`` is +defined as ``log(exp(x[i])/sum(exp(x)))``. + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True then +the corresponding element in :attr:`input` tensor will be included in +log_softmax computation, otherwise the element is ignored. + +The values of masked-out elements of the output tensor have undefined +value: it may or may not be set to zero or nan; the choice may correspond to +the value that leads to the most efficient storage of :attr:`output` +tensor. + +The mask of the log_softmax output tensor can be computed as +``torch.broadcast_to(mask, input.shape)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int): the dimension along which log_softmax is computed. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]]) + >>> input + tensor([[-3., -2., -1.], + [ 0., 1., 2.]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.log_softmax(input, 1, mask=mask) + tensor([[-2.1269, -inf, -0.1269], + [ nan, nan, nan]]) +""" + +logsumexp_docstring = """logsumexp(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor + +Returns logsumexp of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. + +The identity value of logsumexp operation, which is used to start the reduction, is ``-2147483648``. + +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in logsumexp computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of logsumexp operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.logsumexp(input, 1, mask=mask) + tensor([ 0, -9223372036854775808]) +""" + +mean_docstring = """mean(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor + +Returns mean of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. + +By definition, the identity value of a mean operation is the mean +value of the tensor. If all elements of the input tensor along given +dimension(s) :attr:`dim` are masked-out, the identity value of the +mean is undefined. Due to this ambiguity, the elements of output +tensor with strided layout, that correspond to fully masked-out +elements, have ``nan`` values. + +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in mean computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of mean operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.mean(input, 1, mask=mask) + tensor([-2., nan]) +""" + +median_docstring = """median(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor +Returns median of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. +By definition, the identity value of a median operation is the median +value of the tensor. If all elements of the input tensor along given +dimension(s) :attr:`dim` are masked-out, the identity value of the +median is undefined. Due to this ambiguity, the elements of output +tensor with strided layout, that correspond to fully masked-out +elements, have ``nan`` values. +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in median computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of median operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int): the dimension along which median is computed. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. +Example:: + + >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]]) + >>> input + tensor([[-3., -2., -1.], + [ 0., 1., 2.]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.median(input, 1, mask=mask) + tensor([-3., nan]) +""" + +norm_docstring = """norm(input, ord, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor + +Returns norm of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. + +The identity value of norm operation, which is used to start the +reduction, is ``0.0``, except for ``ord=-inf`` it is +``inf``. + +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in norm computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of norm operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + ord (int, float, optional): the order of vector norm. Default: 2. + See :func:`torch.linalg.vector_norm` for a list of supported norms. + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]]) + >>> input + tensor([[-3., -2., -1.], + [ 0., 1., 2.]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.norm(input, 2.0, 1, mask=mask) + tensor([3.1623, 0.0000]) +""" + +normalize_docstring = """normalize(input, ord, dim, *, eps=1e-12, dtype=None, mask=None) -> Tensor + +Returns normalize of all the slices in the :attr:`input` tensor +along :attr:`dim` while the :attr:`input` elements are masked out +according to the boolean tensor :attr:`mask`. + +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Normalize of i-th element in ``x`` is +defined as ``x[i]/max(norm(x, p), eps)``. + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True then +the corresponding element in :attr:`input` tensor will be included in +normalize computation, otherwise the element is ignored. + +The values of masked-out elements of the output tensor have undefined +value: it may or may not be set to zero or nan; the choice may correspond to +the value that leads to the most efficient storage of :attr:`output` +tensor. + +The mask of the normalize output tensor can be computed as +``torch.broadcast_to(mask, input.shape)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + ord (int, float): the order of vector norm. Default: 2. + See :func:`torch.linalg.vector_norm` for a list of supported norms. + dim (int): the dimension along which normalize is computed. + +Keyword args: + eps (float, optional): small value to avoid division by zero. Default: 1e-12. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]]) + >>> input + tensor([[-3., -2., -1.], + [ 0., 1., 2.]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.normalize(input, 2.0, 1, mask=mask) + tensor([[-0.9487, 0.0000, -0.3162], + [ 0.0000, 0.0000, 0.0000]]) +""" + +prod_docstring = """prod(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor + +Returns product of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. + +The identity value of product operation, which is used to start the reduction, is ``1``. + +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in product computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of product operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.prod(input, 1, mask=mask) + tensor([3, 1]) +""" + +softmax_docstring = """softmax(input, dim, *, dtype=None, mask=None) -> Tensor + +Returns softmax of all the slices in the :attr:`input` tensor +along :attr:`dim` while the :attr:`input` elements are masked out +according to the boolean tensor :attr:`mask`. + +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Softmax of i-th element in ``x`` is +defined as ``exp(x[i])/sum(exp(x))``. + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True then +the corresponding element in :attr:`input` tensor will be included in +softmax computation, otherwise the element is ignored. + +The values of masked-out elements of the output tensor have undefined +value: it may or may not be set to zero or nan; the choice may correspond to +the value that leads to the most efficient storage of :attr:`output` +tensor. + +The mask of the softmax output tensor can be computed as +``torch.broadcast_to(mask, input.shape)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int): the dimension along which softmax is computed. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]]) + >>> input + tensor([[-3., -2., -1.], + [ 0., 1., 2.]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.softmax(input, 1, mask=mask) + tensor([[0.1192, 0.0000, 0.8808], + [ nan, nan, nan]]) +""" + +softmin_docstring = """softmin(input, dim, *, dtype=None, mask=None) -> Tensor + +Returns softmin of all the slices in the :attr:`input` tensor +along :attr:`dim` while the :attr:`input` elements are masked out +according to the boolean tensor :attr:`mask`. + +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Softmin of i-th element in ``x`` is +defined as ``exp(-x[i])/sum(exp(-x))``. + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True then +the corresponding element in :attr:`input` tensor will be included in +softmin computation, otherwise the element is ignored. + +The values of masked-out elements of the output tensor have undefined +value: it may or may not be set to zero or nan; the choice may correspond to +the value that leads to the most efficient storage of :attr:`output` +tensor. + +The mask of the softmin output tensor can be computed as +``torch.broadcast_to(mask, input.shape)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int): the dimension along which softmin is computed. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]]) + >>> input + tensor([[-3., -2., -1.], + [ 0., 1., 2.]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.softmin(input, 1, mask=mask) + tensor([[0.8808, 0.0000, 0.1192], + [ nan, nan, nan]]) +""" + +std_docstring = """std(input, dim, unbiased, *, keepdim=False, dtype=None, mask=None) -> Tensor +Returns standard_deviation of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. +The identity value of sample standard deviation operation is undefined. The +elements of output tensor with strided layout, that correspond to +fully masked-out elements, have ``nan`` values. +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in standard_deviation computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of standard_deviation operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``. + unbiased (bool): when True, use Bessel's correction, otherwise, compute + the uncorrected sample variance. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.std(input, 1, False, mask=mask) + tensor([1., nan]) +""" + +sum_docstring = """sum(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor + +Returns sum of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. + +The identity value of sum operation, which is used to start the reduction, is ``0``. + +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in sum computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of sum operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.sum(input, 1, mask=mask) + tensor([-4, 0]) +""" + +var_docstring = """var(input, dim, unbiased, *, keepdim=False, dtype=None, mask=None) -> Tensor +Returns variance of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`. +The identity value of sample variance operation is undefined. The +elements of output tensor with strided layout, that correspond to +fully masked-out elements, have ``nan`` values. +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in variance computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of variance operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``. + unbiased (bool): when True, use Bessel's correction, otherwise, compute + the uncorrected sample variance. + +Keyword args: + keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: False. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: None. + mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. +Example:: + + >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]]) + >>> input + tensor([[-3, -2, -1], + [ 0, 1, 2]]) + >>> mask = tensor([[ True, False, True], [False, False, False]]) + >>> mask + tensor([[ True, False, True], + [False, False, False]]) + >>> torch.masked._ops.var(input, 1, False, mask=mask) + tensor([1., nan]) +""" diff --git a/lib/python3.12/site-packages/torch/masked/_ops.py b/lib/python3.12/site-packages/torch/masked/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..db808c01313306ab29a10ae7453cb2ee7c9e6ea4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/_ops.py @@ -0,0 +1,1798 @@ +# mypy: allow-untyped-defs +import warnings +from typing import Any, Callable, List, Optional, Tuple, TYPE_CHECKING, TypeVar, Union +from typing_extensions import ParamSpec + +import torch +from torch import sym_float, Tensor +from torch._prims_common import corresponding_real_dtype +from torch.masked import _docs +from torch.masked.maskedtensor.core import is_masked_tensor, MaskedTensor +from torch.masked.maskedtensor.creation import as_masked_tensor + + +if TYPE_CHECKING: + from torch.types import _dtype as DType + + DimOrDims = Optional[Union[int, Tuple[int], List[int]]] +else: + # The JIT doesn't understand Union, nor torch.dtype here + DType = int + DimOrDims = Optional[Tuple[int]] + + +__all__: List[str] = [] + +_T = TypeVar("_T") +_P = ParamSpec("_P") + +# All masked reduction/normalization operations have the same +# signatures. Here we introduce docstring templates that are applied +# to docstrings of reduction/normalization functions via +# _apply_docstring_templates decorator. + + +def _apply_docstring_templates(func: Callable[_P, _T]) -> Callable[_P, _T]: + """Decorator that applies docstring templates to function docstring + and returns the function instance. + """ + + doc_string = getattr(_docs, f"{func.__name__}_docstring", None) + if doc_string is None: + warnings.warn( + f"No documentation string available for {func.__name__}." + " PyTorch team should run `python tools/update_masked_docs.py`" + " to generate the missing docstrings." + ) + else: + func.__doc__ = doc_string + + # Expose function as public symbol + __all__.append(func.__name__) + + return func + + +def _generate_docstring(func): + """A utility function called from tools/update_masked_docs.py + script to update the module torch.masked._docs.py + """ + docstring_templates = dict( + reduction_signature="""\ +{function_name}(input, {operation_args}, *, {operation_kwargs}) -> Tensor""", + reduction_descr="""\ +Returns {operation name} of all the elements in the :attr:`input` +tensor along the given dimension(s) :attr:`dim` while the :attr:`input` +elements are masked out according to the boolean tensor +:attr:`mask`.""", + reduction_args="""\ +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of +size 1. Otherwise, :attr:`dim` is squeezed (see +:func:`torch.squeeze`), resulting in the output tensor having 1 (or +``len(dim)``) fewer dimension(s). + +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True +then the corresponding element in :attr:`input` tensor will be +included in {operation name} computation, otherwise the element is +ignored. + +When all elements of :attr:`input` along the given dimension +:attr:`dim` are ignored (fully masked-out), the corresponding element +of the output tensor will have undefined value: it may or may not +correspond to the identity value of {operation name} operation; the +choice may correspond to the value that leads to the most efficient +storage of :attr:`output` tensor. + +The mask of the output tensor can be computed as +``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim, +dtype=torch.bool)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + {args_declarations} + +Keyword args: + {kwargs_declarations}""", + reduction_example="""\ +Example:: + + >>> input = {example_input} + >>> input + {indent_example_input} + >>> mask = {example_mask} + >>> mask + {indent_example_mask} + >>> {full_function_name}(input, {example_args}, mask=mask) + {indent_example_output} +""", + reduction_identity="""\ +The identity value of {operation name} operation, which is used to start the reduction, is ``{identity_int32}``.""", + reduction_identity_dtype="""\ +The identity value of {operation name} operation, which is used to start the +reduction, depends on input dtype. For instance, for float32, uint8, +and int32 dtypes, the identity values are ``{identity_float32}``, ``{identity_uint8}``, and ``{identity_int32}``, respectively.""", + normalization_signature="""\ +{function_name}(input, {operation_args}, *, {operation_kwargs}) -> Tensor""", + normalization_descr="""\ +Returns {operation name} of all the slices in the :attr:`input` tensor +along :attr:`dim` while the :attr:`input` elements are masked out +according to the boolean tensor :attr:`mask`. + +{definition}""", + normalization_args="""\ +The boolean tensor :attr:`mask` defines the "validity" of +:attr:`input` tensor elements: if :attr:`mask` element is True then +the corresponding element in :attr:`input` tensor will be included in +{operation name} computation, otherwise the element is ignored. + +The values of masked-out elements of the output tensor have undefined +value: it may or may not be set to zero or nan; the choice may correspond to +the value that leads to the most efficient storage of :attr:`output` +tensor. + +The mask of the {operation name} output tensor can be computed as +``torch.broadcast_to(mask, input.shape)``. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor +don't need to match, but they must be :ref:`broadcastable +` and the dimensionality of the :attr:`mask` +tensor must not be greater than of the :attr:`input` tensor. + +Args: + input (Tensor): the input tensor + {args_declarations} + +Keyword args: + {kwargs_declarations}""", + normalization_example="""\ +Example:: + + >>> input = {example_input} + >>> input + {indent_example_input} + >>> mask = {example_mask} + >>> mask + {indent_example_mask} + >>> {full_function_name}(input, {example_args}, mask=mask) + {indent_example_output} +""", + ) + + args_and_kwargs = dict( + # argument name sufficies separated by double underscore will + # be removed in the final documentation string. + sum=(("dim",), ("keepdim=False", "dtype=None", "mask=None")), + prod=(("dim",), ("keepdim=False", "dtype=None", "mask=None")), + cumsum=(("dim__as_int",), ("dtype=None", "mask=None")), + cumprod=(("dim__as_int",), ("dtype=None", "mask=None")), + amin=(("dim",), ("keepdim=False", "dtype=None", "mask=None")), + amax=(("dim",), ("keepdim=False", "dtype=None", "mask=None")), + argmin=(("dim__as_int",), ("keepdim=False", "dtype=None", "mask=None")), + argmax=(("dim__as_int",), ("keepdim=False", "dtype=None", "mask=None")), + mean=(("dim",), ("keepdim=False", "dtype=None", "mask=None")), + median=(("dim__as_int",), ("keepdim=False", "dtype=None", "mask=None")), + norm=( + ( + "ord", + "dim", + ), + ("keepdim=False", "dtype=None", "mask=None"), + ), + var=(("dim", "unbiased"), ("keepdim=False", "dtype=None", "mask=None")), + std=(("dim", "unbiased"), ("keepdim=False", "dtype=None", "mask=None")), + logsumexp=(("dim",), ("keepdim=False", "dtype=None", "mask=None")), + softmax=(("dim__as_int",), ("dtype=None", "mask=None")), + log_softmax=(("dim__as_int",), ("dtype=None", "mask=None")), + softmin=(("dim__as_int",), ("dtype=None", "mask=None")), + normalize=( + ( + "ord__required", + "dim__as_int", + ), + ("eps=1e-12", "dtype=None", "mask=None"), + ), + ) + + argument_declarations = dict( + dim="""\ +dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + Default: None that is equivalent to ``tuple(range(input.ndim))``.""", + dim__as_int="""\ +dim (int): the dimension along which {operation name} is computed.""", + ord="""\ +ord (int, float, optional): the order of vector norm. Default: 2. + See :func:`torch.linalg.vector_norm` for a list of supported norms.""", + ord__required="""\ +ord (int, float): the order of vector norm. Default: 2. + See :func:`torch.linalg.vector_norm` for a list of supported norms.""", + unbiased="""\ +unbiased (bool): when True, use Bessel's correction, otherwise, compute + the uncorrected sample variance.""", + eps="""\ +eps (float, optional): small value to avoid division by zero. Default: {default}.""", + keepdim="""\ +keepdim (bool, optional): whether the output tensor has + :attr:`dim` retained or not. Default: {default}.""", + dtype="""\ +dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. Default: {default}.""", + mask="""\ +mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of input tensor + elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.""", + ) + + definitions = dict( + softmax="""\ +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Softmax of i-th element in ``x`` is +defined as ``exp(x[i])/sum(exp(x))``.""", + log_softmax="""\ +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. LogSoftmax of i-th element in ``x`` is +defined as ``log(exp(x[i])/sum(exp(x)))``.""", + softmin="""\ +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Softmin of i-th element in ``x`` is +defined as ``exp(-x[i])/sum(exp(-x))``.""", + normalize="""\ +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Normalize of i-th element in ``x`` is +defined as ``x[i]/max(norm(x, p), eps)``.""", + cumsum="""\ +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Cumsum of i-th element in ``x`` is +defined as ``sum(x[:i])``.""", + cumprod="""\ +Let ``x`` be a sequence of unmasked elements of one-dimensional slice +of the :attr:`input` tensor. Cumsum of i-th element in ``x`` is +defined as ``prod(x[:i])``.""", + ) + + reduction_names = dict( + sum="sum", + prod="product", + amax="maximum", + amin="minimum", + argmax="argmax", + argmin="argmin", + mean="mean", + median="median", + norm="norm", + var="variance", + std="standard_deviation", + logsumexp="logsumexp", + ) + + normalization_names = dict( + softmax="softmax", + log_softmax="log_softmax", + softmin="softmin", + normalize="normalize", + cumsum="cumulative_sum", + cumprod="cumulative_prod", + ) + + operation_names = {} + operation_names.update(reduction_names) + operation_names.update(normalization_names) + + # Default example data: + example_dim = 1 + example_input = torch.tensor([[-3, -2, -1], [0, 1, 2]]) + example_mask = torch.tensor([[True, False, True], [False, False, False]]) + example_args: Tuple[Any, ...] + if func.__name__ in {"norm", "normalize"}: + example_args = (2.0, example_dim) + example_input = example_input.to(dtype=torch.float32) + elif func.__name__ in {"var", "std"}: + example_args = (example_dim, False) + elif func.__name__ == "median": + example_args = (example_dim,) + example_input = example_input.to(dtype=torch.float32) + else: + example_args = (example_dim,) + + operation_args: Tuple[str, ...] + operation_kwargs: Tuple[str, ...] + operation_args, operation_kwargs = args_and_kwargs[func.__name__] + arg_declarations = [ + "\n ".join( + argument_declarations.get(a, f'{a.split("__", 1)[0]}: TBD.').splitlines() + ) + for a in operation_args + ] + kwarg_declarations = [ + "\n ".join( + argument_declarations.get( + a.split("=", 1)[0], f'{a.split("__", 1)[0]}: TBD.' + ) + .format(default=a.split("=", 1)[1]) + .splitlines() + ) + for a in operation_kwargs + ] + + if func.__name__ in reduction_names: + op_kind = "reduction" + doc_sections = ["signature", "descr", "identity", "args", "example"] + elif func.__name__ in normalization_names: + op_kind = "normalization" + doc_sections = ["signature", "descr", "args", "example"] + example_input = example_input.to(dtype=torch.float32) + else: + assert 0 # add function name to operation names dictionaries + example_output = func(example_input, *example_args, mask=example_mask) + + template_data = { + "function_name": func.__name__, + "full_function_name": func.__module__ + "." + func.__name__, + "operation name": operation_names[func.__name__], + "operation_args": ", ".join(a.split("__", 1)[0] for a in operation_args), + "operation_kwargs": ", ".join(a.split("__", 1)[0] for a in operation_kwargs), + # one-line representation of a tensor: + "example_input": " ".join(str(example_input).split()), + "example_args": ", ".join(map(str, example_args)), + "example_mask": " ".join(str(example_mask).split()), + # multi-line representation of a tensor with indent + "indent_example_input": ("\n ").join(str(example_input).splitlines()), + "indent_example_mask": ("\n ").join(str(example_mask).splitlines()), + "indent_example_output": ("\n ").join(str(example_output).splitlines()), + } + + if func.__name__ in reduction_names: + template_data.update( + identity_uint8=_reduction_identity( + func.__name__, torch.tensor(0, dtype=torch.uint8) + ), + identity_int32=_reduction_identity( + func.__name__, torch.tensor(0, dtype=torch.int32) + ), + identity_float32=_reduction_identity( + func.__name__, torch.tensor(0, dtype=torch.float32) + ), + ) + if func.__name__ == "norm": + template_data.update( + identity_ord_ninf=_reduction_identity( + func.__name__, torch.tensor(0, dtype=torch.float32), float("-inf") + ) + ) + elif func.__name__ in normalization_names: + template_data.update(definition=definitions[func.__name__]) + else: + assert 0 # add function name to operation names dictionaries + template_data.update( + args_declarations=("\n ".join(arg_declarations)).format_map(template_data) + ) + template_data.update( + kwargs_declarations=("\n ".join(kwarg_declarations)).format_map( + template_data + ) + ) + + # Apply function name info to docstring templates: + templates = { + k: v.format_map(template_data) + for k, v in docstring_templates.items() + if k.startswith(op_kind) + } + templates.update( + (k, v.format_map(template_data) if isinstance(v, str) else v) + for k, v in template_data.items() + ) + + # Apply docstring templates to function doctring: + if func.__doc__ is None: + doc_template = "\n\n".join([f"{{{op_kind}_{sec}}}" for sec in doc_sections]) + else: + doc_template = func.__doc__ + return doc_template.format_map(templates) + + +def _reduction_identity(op_name: str, input: Tensor, *args): + """Return identity value as scalar tensor of a reduction operation on + given input, or None, if the identity value cannot be uniquely + defined for the given input. + + The identity value of the operation is defined as the initial + value to reduction operation that has a property ``op(op_identity, + value) == value`` for any value in the domain of the operation. + Or put it another way, including or excluding the identity value in + a list of operands will not change the reduction result. + + See https://github.com/pytorch/rfcs/pull/27 for more information. + + """ + dtype: DType = input.dtype + device = input.device + op_name = op_name.rsplit(".", 1)[-1] # lstrip module name when present + if op_name in {"sum", "cumsum"}: + return torch.tensor(0, dtype=dtype, device=device) + elif op_name in {"prod", "cumprod"}: + return torch.tensor(1, dtype=dtype, device=device) + elif op_name in {"amax", "argmax", "logaddexp"}: + if torch.is_floating_point(input): + return torch.tensor(-torch.inf, dtype=dtype, device=device) + elif torch.is_signed(input) or dtype == torch.uint8: + return torch.tensor(torch.iinfo(dtype).min, dtype=dtype, device=device) + elif op_name in {"logsumexp"}: + if torch.is_floating_point(input): + return torch.tensor(-torch.inf, dtype=dtype, device=device) + elif torch.is_complex(input): + return torch.tensor(-torch.inf + 0j, dtype=dtype, device=device) + elif torch.is_signed(input) or dtype == torch.uint8: + return torch.tensor(torch.iinfo(dtype).min, dtype=dtype, device=device) + elif op_name in {"amin", "argmin"}: + if torch.is_floating_point(input): + return torch.tensor(torch.inf, dtype=dtype, device=device) + elif torch.is_signed(input) or dtype == torch.uint8: + return torch.tensor(torch.iinfo(dtype).max, dtype=dtype, device=device) + elif op_name == "mean": + # Strictly speaking, the identity value of the mean operation + # is the mean of the input. Since the mean value depends on + # the dim argument and it may be a non-scalar tensor, we + # consider the identity value of the mean operation ambiguous. + # Moreover, the mean value of empty input is undefined. + return None + elif op_name == "norm": + ord = args[0] if args else 2 + if ord == float("-inf"): + assert torch.is_floating_point(input), input.dtype + return torch.tensor(torch.inf, dtype=dtype, device=device) + return torch.tensor(0, dtype=dtype, device=device) + elif op_name == "median": + # We use NaN for now because the implementation is currently using torch.nanmedian + # and NaN is the identity for that function since it gets ignored + dtype = input.dtype if torch.is_floating_point(input) else torch.float + return torch.tensor(torch.nan, dtype=dtype, device=device) + elif op_name in {"var", "std"}: + return None + raise NotImplementedError(f"identity of {op_name} on {dtype} input") + + +def _canonical_dim(dim: DimOrDims, ndim: int) -> Tuple[int, ...]: + """Return dim argument as a tuple of sorted dim values.""" + dims: List[int] = [] + if dim == (): + # Currently, `dim=()` in reductions operations means "reduce + # over all dimensions" while in future, it will read "no + # reduce". See https://github.com/pytorch/pytorch/issues/29137 + # When gh-29137 is resolved, this if-block must be deleted. + dim = None + if dim is None: + return tuple(range(ndim)) + ndim = max(ndim, 1) + dim_ = (dim,) if isinstance(dim, (int, torch.SymInt)) else dim + for d in dim_: + if d in dims: + raise RuntimeError(f"dim={d} appears multiple times in the list of dims") + if d >= ndim or d < -ndim: + raise IndexError( + f"Dimension out of range (expected to be in range of [{-ndim}, {ndim - 1}], but got {d})" + ) + dims.append(d % ndim) + return tuple(sorted(dims)) + + +def _sparse_coo_flatten_indices(indices: Tensor, shape: tuple): + # Flatted N-D indices to 1-D indices + flat_indices = indices.new_zeros(indices.size(1)) + for d, sz in enumerate(shape): + flat_indices.mul_(sz) + flat_indices.add_(indices[d]) + return flat_indices + + +def _any(input: Tensor, dim: tuple, keepdim: bool): + # Support torch.any with tuple dim argument. + # Workaround of https://github.com/pytorch/pytorch/issues/56586 + r = input + for d in reversed(dim): + r = r.any(dim=d, keepdim=keepdim) + return r + + +def _sparse_coo_where(mask: Tensor, input: Tensor, fill_value: Tensor) -> Tensor: + """Sparse variant of torch.where. Supports sparse COO and hybrid sparse COO tensors. + + _sparse_coo_where implements the following invariant: + + _sparse_coo_where(mask, input, fill_value).to_dense(fill_value) == + torch.where(mask.to_dense(), input.to_dense(), torch.full(input.shape, fill_value)) + + where `a == b` means `assertEqual(a, b)`, mask is boolean sparse + tensor, and `to_dense(fill_value)` is like `to_dense()` except + that the unspecified elements are mapped to `fill_value` rather + than to `0`. + + Returns a sparse COO tensor with the following features: + + - all specified elements correspond to masked-in elements that + have the values of the input tensor. If there exists a masked-in + element (as specified by mask) that is not specified in the + input, in the result tensor, the corresponding element has value + 0. In the dense part of the sparse tensor, the masked-out + elements are replaced with fill_value. + + - all unspecified elements correspond to masked-out elements. + """ + + assert input.layout == torch.sparse_coo + assert mask.layout == input.layout + assert mask.shape == input.shape + assert mask.dense_dim() == input.dense_dim() # TODO: eliminate this restriction + + input = input.coalesce() + + # For set operations on sparse tensor indices, we'll convert + # multi-dimensional indices to 1-D indices for efficiency. + input_flat_indices = _sparse_coo_flatten_indices( + input.indices(), input.shape[: input.sparse_dim()] + ) + mask_flat_indices = _sparse_coo_flatten_indices( + mask.indices(), mask.shape[: mask.sparse_dim()] + ) + + # the set of mask flat indices that define masked-in elements: + if mask.dense_dim() > 0: + mask_values = _any( + mask.values(), tuple(range(1, input.sparse_dim() + 1)), False + ) + else: + mask_values = mask.values() + maskin_flat_indices = mask_flat_indices[mask_values.nonzero()[:, 0]] + + def intersection(i1, i2): + union, counts = torch.cat([i1, i2]).unique(return_counts=True) + return union, torch.where(counts.gt(1)) + + def minus(i1, i2): + union, counts = torch.cat([i1, i2]).unique(return_counts=True) + return intersection(union[torch.where(counts.eq(1))], i1) + + def _apply(a): + obj, w = a + return obj[w] + + # the set of input flat indices of specified and masked-in elements: + maskin_input_flat_indices = _apply( + intersection(maskin_flat_indices, input_flat_indices) + ) + _, w = intersection(input_flat_indices, maskin_input_flat_indices) + + # the indices and values of masked-in elements + where_input_indices = input.indices()[(slice(None),) + w] + where_input_values = input.values()[w] + + if mask.dense_dim() > 0: + # apply mask to the dense part of the input values: + _, w1 = intersection(mask_flat_indices, maskin_input_flat_indices) + where_mask_values = mask.values()[w1] + where_input_values = torch.where( + where_mask_values, where_input_values, fill_value + ) + + # the set of flat indices of unspecified input and masked-in elements: + maskin_zero_flat_indices = _apply( + minus(maskin_flat_indices, maskin_input_flat_indices) + ) + + # the indices of masked-in zero elements + _, w = intersection(mask_flat_indices, maskin_zero_flat_indices) + where_zero_indices = mask.indices()[(slice(None),) + w] + + # construct result + n = where_zero_indices.size(1) + if n == 0: + # the input is coalesced, hence input_flat_indices are ordered + # and the result is guaranteed to be coalesced: + result = torch.sparse_coo_tensor( + where_input_indices, where_input_values, input.shape + ) + return result._coalesced_(True) + + where_indices = torch.cat([where_input_indices, where_zero_indices], dim=1) + where_values = torch.cat( + [ + where_input_values, + where_input_values.new_zeros((n,) + where_input_values.shape[1:]), + ] + ) + result = torch.sparse_coo_tensor(where_indices, where_values, input.shape) + + # appending zero elements leads to uncoalesced sparse tensor + return result.coalesce() + + +def _sparse_coo_scatter_reduction_helper( + op, + mask_input: Tensor, + dims: Tuple[int, ...], + keepdim: bool, + dtype: Optional[DType] = None, +) -> Tensor: + reduce = op.__name__ + valid_reductions = ["sum", "prod", "amax", "amin"] + if reduce not in valid_reductions: + raise ValueError( + f"op must be one of {' '.join(valid_reductions)}, but got {reduce} instead" + ) + + output_dtype = dtype + values, indices = mask_input._values(), mask_input._indices() + input_dims = mask_input.dim() + num_sparse_dims = mask_input.sparse_dim() + reduced_sparse_dims = [] + retained_sparse_dims = [] + reduced_dense_dims = [] + + # promote dtype if specified + if values.dtype != output_dtype: + values = values.to(output_dtype) + + if keepdim: + output_shape = tuple( + 1 if i in dims else si for (i, si) in enumerate(mask_input.shape) + ) + else: + output_shape = tuple( + si for (i, si) in enumerate(mask_input.shape) if i not in dims + ) + + for d in dims: + if d >= input_dims: + continue + + if d < num_sparse_dims: + reduced_sparse_dims.append(d) + else: + reduced_dense_dims.append(d + 1 - num_sparse_dims) + + # Reduce dense dimensions + if len(reduced_dense_dims) > 0: + if reduce == "sum": + new_values = values + new_values = op(new_values, dim=reduced_dense_dims, keepdim=bool(keepdim)) + else: + # FIXME: Implement reductions for dense dimensions for ops with non-zero reduction identities + return NotImplemented + else: + new_values = values.clone() + + # Reduce sparse dimensions + if len(reduced_sparse_dims) == num_sparse_dims: + if reduce in {"amax", "amin"} and new_values.size(0) == 0: + # IndexError: amax(): Expected reduction dim 0 to have non-zero size. + # sum()/prod() return the reduction identity when dim has size 0 but amax()/amin() do not + # See https://github.com/pytorch/pytorch/issues/61901 + new_values = _reduction_identity(reduce, new_values) + else: + new_values = op(new_values, dim=0) + if keepdim: + for _ in range(num_sparse_dims): + new_values = new_values.unsqueeze(0) + return new_values.to(dtype=output_dtype).to_sparse() + else: + new_indices = indices.clone() + if keepdim: + # zero out reduced sparse dimensions if keepdim = True + # ensures that the call to torch.unique folds duplicated indices together while preserving the dimension + new_indices[reduced_sparse_dims, :] = 0 + else: + # remove reduced sparse dimensions if keepdim = False + if len(reduced_sparse_dims) > 0: + retained_sparse_dims = [ + i + for i in range(num_sparse_dims) + if i not in set(reduced_sparse_dims) + ] + new_indices = new_indices.index_select( + 0, torch.tensor(retained_sparse_dims).to(mask_input.device) + ) + + # Use scatter_reduce to reduce items in the new_values tensor that correspond to the same indices in new_indices + if new_indices.numel() > 0: + # lexsort indices and get index tensor for scatter reduction + new_indices, inverse_indices = torch.unique( + new_indices, return_inverse=True, dim=1 + ) + out_shape = list(new_values.shape) + out_shape[0] = new_indices.shape[1] + for _ in range(new_values.ndim - 1): + inverse_indices = inverse_indices.unsqueeze(-1) + scatter_indices = inverse_indices.expand(new_values.shape) + # FIXME: temporary workaround for issue with bfloat16/float16 remove when acctype is implemented for scatter_reduce + if output_dtype in {torch.bfloat16, torch.float16}: + new_values = new_values.to(torch.float) + out = new_values.new_empty(out_shape) + new_values = out.scatter_reduce_( + 0, scatter_indices, new_values, reduce=reduce, include_self=False + ) + new_values = new_values.to(dtype=output_dtype) + else: + out = new_values.new_empty(out_shape) + new_values = out.scatter_reduce_( + 0, scatter_indices, new_values, reduce=reduce, include_self=False + ) + + return torch.sparse_coo_tensor( + new_indices, + new_values, + output_shape, + dtype=output_dtype, + device=mask_input.device, + ) + + +def _sparse_csr_segment_reduction_helper( + op, + mask_input: Tensor, + dims: Tuple[int, ...], + keepdim: bool, + dtype: Optional[DType] = None, +) -> Tensor: + # Currently, while sparse CSR is always 2D with no dense dimensions keepdim must be True + # FIXME: when dense dimensions are implemented for CSR tensors + assert ( + keepdim + ), "reduction operations on CSR tensors with keepdim=False is unsupported" + reduce = op.__name__ + valid_reductions = ["sum", "prod", "mean", "amax", "amin"] + if reduce not in valid_reductions: + raise ValueError( + f"op must be one of {' '.join(valid_reductions)}, but got {reduce} instead" + ) + device = mask_input.device + output_dtype = dtype + values, crow_indices, col_indices = ( + mask_input.values(), + mask_input.crow_indices(), + mask_input.col_indices(), + ) + + # promote dtype if specified + if values.dtype != output_dtype: + values = values.to(output_dtype) + + if len(dims) == 0: + return mask_input + if len(dims) == 1: + if dims[0] == 0: + new_col_indices, scatter_indices = torch.unique( + col_indices, return_inverse=True + ) + new_nnz = new_col_indices.shape[0] + new_crow_indices = torch.tensor([0, new_nnz]) + new_values = values.new_empty(new_col_indices.shape) + new_values.scatter_reduce_( + 0, scatter_indices, values, reduce, include_self=False + ) + new_shape = [1, mask_input.size(1)] + else: + assert ( + dims[0] == 1 + ), "Sparse CSR tensors are 2D and only support reduction along dim 0 or 1." + # all intervals new_crow_indices[i] - new_crow_indices[i-1] are 1 + # except for where crow_indices[i] == crow_indices[i-1] where the interval remains as 0 + new_crow_indices = torch.cat( + ( + crow_indices.new_zeros(1), + torch.cumsum(torch.diff(crow_indices) != 0, 0), + ), + 0, + ) + new_nnz = new_crow_indices[-1] + new_col_indices = col_indices.new_zeros(new_nnz) + new_values = torch._segment_reduce(values, reduce, offsets=crow_indices) # type: ignore[attr-defined] + new_shape = [mask_input.size(0), 1] + else: + assert len(dims) == 2 + nnz = min(1, values.numel()) + if nnz == 1: + op_kwargs = {"keepdim": True, "dtype": output_dtype} + # amax and amin do not support dtype kwarg + if reduce in ["amax", "amin"]: + del op_kwargs["dtype"] + new_values = op(values, 0, **op_kwargs) + else: + new_values = torch.empty(0, dtype=output_dtype) + new_col_indices = col_indices.new_zeros(nnz) + new_crow_indices = torch.tensor([0, nnz]) + new_shape = [1, nnz] + + return torch.sparse_csr_tensor( + new_crow_indices, + new_col_indices, + new_values, + new_shape, + dtype=output_dtype, + device=device, + ) + + +def _sparse_csr_where(mask: Tensor, input: Tensor, fill_value: Tensor) -> Tensor: + """Sparse variant of torch.where. Supports sparse CSR tensors.""" + # TODO: implement sparse CSR specific where operator for efficiency + return _sparse_coo_where( + mask.to_sparse_coo(), input.to_sparse_coo(), fill_value + ).to_sparse_csr() + + +def _where(mask: Tensor, input: Tensor, fill_value: Tensor) -> Tensor: + """torch.where with sparse inputs support. + + _where implements the following invariant: + + _where(mask, input, fill_value).to_dense(fill_value) == + torch.where(mask.to_dense(), input.to_dense(), torch.full(input.shape, fill_value)) + + where `a == b` means `assertEqual(a, b)`, mask is boolean sparse + tensor, and `to_dense(fill_value)` is like `to_dense()` except + that the unspecified elements are mapped to `fill_value` rather + than to `0`. + + Returns a sparse tensor with the following features: + + - all specified elements correspond to masked-in elements that + have the values of the input tensor. If there exists a masked-in + element (as specified by mask) that is not specified in the + input, in the result tensor, the corresponding element has value + 0. In the dense part of the sparse tensor, the masked-out + elements are replaced with fill_value. + + - all unspecified elements correspond to masked-out elements. + """ + if mask.layout == torch.strided: + return torch.where(mask, input, fill_value) + elif mask.layout == torch.sparse_coo: + return _sparse_coo_where(mask, input, fill_value) + elif mask.layout == torch.sparse_csr: + return _sparse_csr_where(mask, input, fill_value) + else: + raise ValueError( + f"_where expects strided or sparse COO or sparse CSR tensor but got {mask.layout}" + ) + + +def _input_mask(input: Union[Tensor, MaskedTensor], *args, **kwargs) -> Tensor: + """Return canonical input mask. + + A canonical input mask is defined as a boolean mask tensor that + shape and layout matches with the shape and the layout of the + input. + + The canonical input mask is computed from the :attr:`mask` tensor + content to meet the following criteria: + + 1. The shape of the canonical input mask is the same as the shape + of :attr:`input` tensor. If the mask tensor has a smaller shape + than the shape of the :attr:`input`, broadcasting rules will be + applied. Downcasting of mask is not supported. + + 2. The layout of the canonical input mask is the same as the + layout of the :attr:`input` tensor. If the mask has different + layout, it will be converted to the expected layout. In the + case of sparse COO layout, the canonical input mask will be + coalesced. + + 3. The dtype of the canonical input mask is torch.bool. If the + mask dtype is not bool then it will be converted to bool dtype + using `.to(dtype=bool)` method call. + + 4. The elements of the canonical input mask have boolean values + copied from the content of the :attr:`mask` tensor (after + possible broadcasting and dtype conversion transforms). In + general, the sparsity pattern of the sparse canonical input + mask need not to be the same as the sparsity pattern of the + sparse :attr:`input` tensor. + + """ + if input.layout not in {torch.strided, torch.sparse_coo, torch.sparse_csr}: + raise ValueError( + f"_input_mask expects strided or sparse COO or sparse CSR tensor but got {input.layout}" + ) + + mask = kwargs.get("mask") + + # default mask + if mask is None: + raise ValueError("_input_mask requires explicit mask") + + # mask shape must match with input shape + if mask.shape != input.shape: + if mask.ndim > input.ndim: + raise IndexError( + "_input_mask expected broadcastable mask (got mask dimensionality higher than of the input)" + ) + if mask.layout == torch.strided: + mask = torch.broadcast_to(mask.clone(), input.shape).to(dtype=torch.bool) + elif mask.layout == torch.sparse_coo: + mask = torch._sparse_broadcast_to(mask, input.shape) + else: + assert mask.layout == torch.sparse_csr + # Broadcasting of CSR tensors is not implemented. Working + # around by using COO layout. + mask = torch._sparse_broadcast_to( + mask.to_sparse(), input.shape + ).to_sparse_csr() + + # mask layout must match with input layout + if mask.layout != input.layout: + if input.layout == torch.strided: + mask = mask.to_dense() + elif input.layout == torch.sparse_coo: + if mask.layout == torch.strided: + mask = mask.to_sparse(input.sparse_dim()) + else: + mask = mask.to_sparse() + else: + assert input.layout == torch.sparse_csr + mask = mask.to_sparse_csr() + + # sparse mask must be coalesced + if mask.layout == torch.sparse_coo: + mask = mask.coalesce() + + # mask is a boolean tensor + mask = mask.to(dtype=torch.bool) + + return mask + + +def _output_mask(op, input: Tensor, *args, **kwargs) -> Tensor: + """Return output mask of masked operation applied to given arguments.""" + if callable(op): + is_reduction = op.__name__ in { + "sum", + "prod", + "amax", + "amin", + "argmax", + "argmin", + "mean", + "median", + "norm", + "var", + "std", + "logsumexp", + } + is_normalization = op.__name__ in { + "softmax", + "log_softmax", + "softmin", + "normalize", + "cumsum", + "cumprod", + } + if is_reduction: + if op.__name__ == "norm": + if args: + args = args[1:] # lstrip ord argument + dim = args[0] if args else kwargs.get("dim") + outmask = _input_mask(input, *args, **kwargs) + keepdim = kwargs.get("keepdim", False) + dim_ = _canonical_dim(dim, input.ndim) + return _any(outmask, dim_, bool(keepdim)) + elif is_normalization: + return _input_mask(input, *args, **kwargs) + else: + raise ValueError( + f"_output_mask expected masked operation (got callable {op.__module__}.{op.__name__})" + ) + else: + raise ValueError( + f"_output_mask expected masked operation (got {type(op).__name__} object)" + ) + + +def _combine_input_and_mask( + op, input: Union[MaskedTensor, Tensor], mask, *args +) -> Tensor: + def helper(input, mask): + if mask is None: + return input + canonical_mask = _input_mask(input, mask=mask) + if callable(op): + fill_value = _reduction_identity(op.__name__, input, *args) + return _where(canonical_mask, input, fill_value) + else: + raise ValueError( + f"_combine_input_and_mask expected masked operation (got {type(op).__name__} object)" + ) + + class Combine(torch.autograd.Function): + @staticmethod + def forward(ctx, input, mask): + """Return input with masked-out elements eliminated for the given operations.""" + ctx.save_for_backward(mask) + + if mask is not None: + ctx.mark_non_differentiable(mask) + + return helper(input, mask) + + @staticmethod + def backward(ctx, grad_output): + (mask,) = ctx.saved_tensors + grad_data = ( + grad_output.get_data() if is_masked_tensor(grad_output) else grad_output + ) + result = as_masked_tensor(grad_data, mask) + return result, None + + return ( + Combine.apply(input.get_data(), input.get_mask()) # type: ignore[union-attr] + if is_masked_tensor(input) + else helper(input, mask) + ) + + +@_apply_docstring_templates +def sum( + input: Union[Tensor, MaskedTensor], + dim: DimOrDims = None, + *, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + # __doc__ is generated by _apply_docstring_templates decorator + if dtype is None: + # promote integer types to int64 when output dtype is not specified + if input.layout == torch.sparse_csr: + if input.dtype in { + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + torch.int32, + }: + # csr.to(dtype=torch.int64) is not implemented, so + # using coo.to on input to ensure the promoted dtype + input = input.to_sparse_coo().to(dtype=torch.int64).to_sparse_csr() + else: + dtype = input.dtype + else: + dtype = input.dtype + if input.dtype in { + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + torch.int32, + }: + dtype = torch.int64 + dim_ = _canonical_dim(dim, input.ndim) + mask_input = _combine_input_and_mask(sum, input, mask) + if mask_input.layout == torch.strided: + return torch.sum(mask_input, dim_, bool(keepdim), dtype=dtype) + elif mask_input.layout == torch.sparse_coo: + return _sparse_coo_scatter_reduction_helper( + torch.sum, mask_input, dim_, bool(keepdim), dtype + ) + elif mask_input.layout == torch.sparse_csr: + return torch._sparse_csr_sum( + mask_input, dim=list(dim_), keepdim=bool(keepdim), dtype=dtype + ) + else: + raise ValueError( + f"masked sum expects strided, sparse_coo or sparse_csr tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def prod( + input: Union[Tensor, MaskedTensor], + dim: DimOrDims = None, + *, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + # __doc__ is generated by _apply_docstring_templates decorator + if dtype is None: + # promote integer types to int64 when output dtype is not specified + if input.layout == torch.sparse_csr: + if input.dtype in { + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + torch.int32, + }: + # csr.to(dtype=torch.int64) is not implemented, so + # using coo.to on input to ensure the promoted dtype + input = input.to_sparse_coo().to(dtype=torch.int64).to_sparse_csr() + else: + dtype = input.dtype + else: + dtype = input.dtype + if input.dtype in { + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + torch.int32, + }: + dtype = torch.int64 + dim_ = _canonical_dim(dim, input.ndim) + mask_input = _combine_input_and_mask(prod, input, mask) + if mask_input.layout == torch.strided: + # Workaround https://github.com/pytorch/pytorch/issues/56586 + result = mask_input + result = result.to(dtype=dtype) + for d in reversed(dim_): + result = result.prod(dim=d, keepdim=bool(keepdim)) + return result + elif mask_input.layout == torch.sparse_coo: + if mask is None: + # See comment in the sparse_csr branch, the same issue arises for sparse_coo tensors + raise ValueError( + "masked prod expects explicit mask for sparse_coo tensor input" + ) + return _sparse_coo_scatter_reduction_helper( + torch.prod, mask_input, dim_, bool(keepdim), dtype + ) + elif mask_input.layout == torch.sparse_csr: + if mask is None: + # mask is None corresponds to all-True mask. The + # unspecified elements in the CSR tensor correspond to + # zero values. Hence, the prod reduction result is + # automatically zero unless all elements are specified. + # A semi-optimal way to take this into account is to use: + # + # masked_prod(csr, ..., mask=None) == torch._sparse_csr_prod(csr, ...) * all(csr.nonzero(), ...) + # + # but that requires implementing `all` and `nonzero` + # support for sparse csr tensors. + raise ValueError( + "masked prod expects explicit mask for sparse_csr tensor input" + ) + return torch._sparse_csr_prod( + mask_input, dim=list(dim_), keepdim=bool(keepdim), dtype=dtype + ) + else: + raise ValueError( + f"masked prod expects strided, sparse_coo or sparse_csr tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def cumsum( + input: Tensor, + dim: int, + *, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + if dtype is None: + dtype = input.dtype + dim_ = _canonical_dim(dim, input.ndim)[0] + mask_input = _combine_input_and_mask(sum, input, mask) + if mask_input.layout == torch.strided: + return torch.cumsum(mask_input, dim_, dtype=dtype).to(dtype=dtype) + else: + raise ValueError( + f"masked cumsum expects strided tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def cumprod( + input: Tensor, + dim: int, + *, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + if dtype is None: + dtype = input.dtype + dim_ = _canonical_dim(dim, input.ndim)[0] + mask_input = _combine_input_and_mask(prod, input, mask) + if mask_input.layout == torch.strided: + return torch.cumprod(mask_input, dim_, dtype=dtype).to(dtype=dtype) + else: + raise ValueError( + f"masked cumprod expects strided tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def amax( + input: Union[Tensor, MaskedTensor], + dim: DimOrDims = None, + *, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + """\ +{reduction_signature} + +{reduction_descr} + +{reduction_identity_dtype} + +{reduction_args} + +{reduction_example}""" + if dtype is None: + dtype = input.dtype + + mask_input = _combine_input_and_mask(amax, input, mask) + dim_ = _canonical_dim(dim, mask_input.ndim) + if mask_input.layout == torch.strided: + return torch.amax(mask_input, dim_, bool(keepdim)).to(dtype=dtype) + elif mask_input.layout == torch.sparse_coo: + if mask is None: + # See comment in the sparse_csr branch of prod, a similar issue arises here + # where unspecified elements along a dimension may need to be reduced with the result + raise ValueError( + "masked amax expects explicit mask for sparse_coo tensor input" + ) + return _sparse_coo_scatter_reduction_helper( + torch.amax, mask_input, dim_, bool(keepdim), dtype + ) + elif mask_input.layout == torch.sparse_csr: + if mask is None: + raise ValueError( + "masked amax expects explicit mask for sparse_csr tensor input" + ) + return _sparse_csr_segment_reduction_helper( + torch.amax, mask_input, dim_, bool(keepdim), dtype + ) + else: + raise ValueError( + f"masked amax expects strided, sparse_coo or sparse_csr tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def amin( + input: Union[Tensor, MaskedTensor], + dim: DimOrDims = None, + *, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + """\ +{reduction_signature} + +{reduction_descr} + +{reduction_identity_dtype} + +{reduction_args} + +{reduction_example}""" + if dtype is None: + dtype = input.dtype + + mask_input = _combine_input_and_mask(amin, input, mask) + dim_ = _canonical_dim(dim, mask_input.ndim) + if mask_input.layout == torch.strided: + return torch.amin(mask_input, dim_, bool(keepdim)).to(dtype=dtype) + elif mask_input.layout == torch.sparse_coo: + if mask is None: + # See comment in the sparse_csr branch of prod, a similar issue arises here + # where unspecified elements along a dimension may need to be reduced with the result + raise ValueError( + "masked amax expects explicit mask for sparse_coo tensor input" + ) + return _sparse_coo_scatter_reduction_helper( + torch.amin, mask_input, dim_, bool(keepdim), dtype + ) + elif mask_input.layout == torch.sparse_csr: + if mask is None: + raise ValueError( + "masked amin expects explicit mask for sparse_csr tensor input" + ) + return _sparse_csr_segment_reduction_helper( + torch.amin, mask_input, dim_, bool(keepdim), dtype + ) + else: + raise ValueError( + f"masked amin expects strided, sparse_coo or sparse_csr tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def argmax( + input: Union[Tensor, MaskedTensor], + dim: Optional[int] = None, + *, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + """\ +{reduction_signature} +{reduction_descr} +{reduction_identity_dtype} +{reduction_args} +{reduction_example}""" + if dtype is None: + dtype = input.dtype + mask_input = _combine_input_and_mask(argmax, input, mask) + if mask_input.layout == torch.strided: + return torch.argmax(mask_input, dim, bool(keepdim)).to(dtype=dtype) + else: + raise ValueError( + f"masked argmax expects strided tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def argmin( + input: Union[Tensor, MaskedTensor], + dim: Optional[int] = None, + *, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + """\ +{reduction_signature} +{reduction_descr} +{reduction_identity_dtype} +{reduction_args} +{reduction_example}""" + if dtype is None: + dtype = input.dtype + mask_input = _combine_input_and_mask(argmin, input, mask) + if mask_input.layout == torch.strided: + return torch.argmin(mask_input, dim, bool(keepdim)).to(dtype=dtype) + else: + raise ValueError( + f"masked argmin expects strided tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def mean( + input: Union[Tensor, MaskedTensor], + dim: DimOrDims = None, + *, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + """\ +{reduction_signature} + +{reduction_descr} + +By definition, the identity value of a mean operation is the mean +value of the tensor. If all elements of the input tensor along given +dimension(s) :attr:`dim` are masked-out, the identity value of the +mean is undefined. Due to this ambiguity, the elements of output +tensor with strided layout, that correspond to fully masked-out +elements, have ``nan`` values. + +{reduction_args} + +{reduction_example}""" + if dtype is None: + dtype = input.dtype + if input.layout == torch.strided: + if mask is None: + # TODO: compute count analytically + count = sum( + torch.ones(input.shape, dtype=torch.int64, device=input.device), + dim, + keepdim=keepdim, + ) + total = sum(input, dim, keepdim=keepdim, dtype=dtype) + else: + inmask = _input_mask(input, mask=mask) + count = inmask.sum(dim=dim, keepdim=bool(keepdim)) + total = sum(input, dim, keepdim=keepdim, dtype=dtype, mask=inmask) + return total / count + elif input.layout == torch.sparse_csr: + mask_input = _combine_input_and_mask(mean, input, mask) + dim_ = _canonical_dim(dim, mask_input.ndim) + if mask is None: + raise ValueError( + "masked mean expects explicit mask for sparse_csr tensor input" + ) + return _sparse_csr_segment_reduction_helper( + torch.mean, mask_input, dim_, bool(keepdim), dtype + ) + else: + raise ValueError( + f"masked mean expects strided or sparse_csr tensor (got {input.layout} tensor)" + ) + + +@_apply_docstring_templates +def median( + input: Union[Tensor, MaskedTensor], + dim: int = -1, + *, + keepdim: bool = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + """\ +{reduction_signature} +{reduction_descr} +By definition, the identity value of a median operation is the median +value of the tensor. If all elements of the input tensor along given +dimension(s) :attr:`dim` are masked-out, the identity value of the +median is undefined. Due to this ambiguity, the elements of output +tensor with strided layout, that correspond to fully masked-out +elements, have ``nan`` values. +{reduction_args} +{reduction_example}""" + if dtype is None: + dtype = input.dtype + dim_ = _canonical_dim(dim, input.ndim)[0] + is_float = torch.is_floating_point(input) + if not is_float: + input = input.to(dtype=torch.float) + mask_input = _combine_input_and_mask(median, input, mask) + if mask_input.layout == torch.strided: + output = torch.nanmedian(mask_input, dim_, keepdim).values + if is_float: + return output + elif not is_float and not torch.isnan(output).any(): + return output.to(dtype=dtype) + else: + raise ValueError( + "masked median expects no fully masked out rows if dtype is not floating point" + ) + else: + raise ValueError( + f"masked median expects strided tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def logsumexp( + input: Tensor, + dim: DimOrDims = None, + *, + keepdim: bool = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + if dtype is None: + dtype = input.dtype + dim_ = _canonical_dim(dim, input.ndim) + mask_input = _combine_input_and_mask(logsumexp, input, mask) + if mask_input.layout == torch.strided: + return torch.logsumexp(mask_input, dim_, keepdim=keepdim).to(dtype=dtype) + else: + raise ValueError( + f"masked logsumexp expects strided tensor (got {mask_input.layout} tensor)" + ) + + +# Cannot use _apply_docstring_templates as it is only set up for reductions and normalizations +def logaddexp( + input: Union[Tensor, MaskedTensor], + other: Union[Tensor, MaskedTensor], + *, + dtype: Optional[DType] = None, + input_mask: Optional[Tensor] = None, + other_mask: Optional[Tensor] = None, +) -> Tensor: + """logaddexp(input, other, *, dtype=None, input_mask=None, other_mask=None) -> Tensor + + Returns logaddexp of all the elements in the :attr:`input` and the :attr:`other` + tensor. The :attr:`input` elements are masked out according to the boolean tensor + :attr:`input_mask` and the attr:`other` elements are masked out according to the boolean tensor + :attr:`other_mask`. + + The shapes of a mask tensor and the tensor to be masked + don't need to match, but they must be :ref:`broadcastable + ` and the dimensionality of the mask + tensor must not be greater than of the tensor to be masked. + + Args: + input (Tensor): the input tensor + other (Tensor): the second input tensor + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the output tensor is + casted to :attr:`dtype` after the operation is + performed. Default: None. + input_mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of :attr:`input` tensor elements. + Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``. + other_mask (:class:`torch.Tensor`, optional): the boolean tensor + containing the binary mask of validity of :attr:`other` tensor elements. + Default: None that is equivalent to ``torch.ones(other.shape, dtype=torch.bool)``. + + Example:: + + >>> input = torch.tensor([-100.0, -200, -300]) + >>> input + tensor([-100., -200., -300.]) + >>> other = torch.tensor([-1.0, -2, -3]) + >>> other + tensor([-1., -2., -3.]) + >>> mask = torch.tensor([True, False, True]) + >>> mask + tensor([ True, False, True]) + >>> torch.masked._ops.logaddexp(input, other, input_mask=mask, other_mask=mask) + tensor([-1., -inf, -3.])""" + if dtype is None: + dtype = input.dtype + if input.layout == torch.strided and other.layout == torch.strided: + mask_input = _combine_input_and_mask(logaddexp, input, input_mask) + mask_other = _combine_input_and_mask(logaddexp, other, other_mask) + return torch.logaddexp(mask_input, mask_other).to(dtype=dtype) + else: + raise ValueError( + f"masked logaddexp expects strided tensors (got {input.layout} tensor for input, {other.layout} for other)" + ) + + +@_apply_docstring_templates +def norm( + input: Union[Tensor, MaskedTensor], + ord: Optional[float] = 2.0, + dim: DimOrDims = None, + *, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + """\ +{reduction_signature} + +{reduction_descr} + +The identity value of norm operation, which is used to start the +reduction, is ``{identity_float32}``, except for ``ord=-inf`` it is +``{identity_ord_ninf}``. + +{reduction_args} + +{reduction_example}""" + if dtype is None: + dtype = input.dtype + mask_input = _combine_input_and_mask(norm, input, mask, ord) + if mask_input.layout == torch.strided: + dim_ = _canonical_dim(dim, input.ndim) + return torch.linalg.vector_norm( + mask_input, ord, dim_, bool(keepdim), dtype=dtype + ) + else: + raise ValueError( + f"masked norm expects strided tensor (got {mask_input.layout} tensor)" + ) + + +def _std_var( + input: Union[Tensor, MaskedTensor], + dim: DimOrDims, + unbiased: Optional[bool], + *, + correction_opt: Optional[Union[int, float]], + keepdim: Optional[bool], + dtype: Optional[DType], + mask: Optional[Tensor], + take_sqrt: Optional[bool], +) -> Tensor: + assert ( + unbiased is None or correction_opt is None + ), "Only one of unbiased and correction may be given" + correction = 1.0 + if unbiased is not None: + correction = 1.0 if unbiased else 0.0 + if correction_opt is not None: + correction = sym_float(correction_opt) + + if dtype is None: + dtype = input.dtype + if not (dtype.is_floating_point or dtype.is_complex): + dtype = torch.float32 + compute_dtype = dtype + if not (compute_dtype.is_floating_point or compute_dtype.is_complex): + compute_dtype = torch.float32 + if input.layout == torch.strided: + if mask is None: + # TODO: compute count analytically + count = sum( + torch.ones(input.shape, dtype=torch.int64, device=input.device), + dim, + keepdim=True, + ) + sample_total = sum(input, dim, keepdim=True, dtype=dtype) + else: + inmask = _input_mask(input, mask=mask) + count = inmask.sum(dim=dim, keepdim=True) + sample_total = sum(input, dim, keepdim=True, dtype=dtype, mask=inmask) + # TODO: replace torch.subtract/divide/square/maximum with + # masked subtract/divide/square/maximum when these will be + # available. + sample_mean = torch.divide(sample_total, count) + x = torch.subtract(input, sample_mean) + if mask is None: + total = sum(x * x.conj(), dim, keepdim=keepdim, dtype=compute_dtype) + else: + total = sum( + x * x.conj(), dim, keepdim=keepdim, dtype=compute_dtype, mask=inmask # type: ignore[possibly-undefined] + ) + if not keepdim: + count = count.reshape(total.shape) + if correction != 0: + real_dtype = ( + corresponding_real_dtype(compute_dtype) + if compute_dtype.is_complex + else compute_dtype + ) + count = count.to(real_dtype) + count = torch.subtract(count, correction) + count = torch.maximum(count, count.new_zeros([])) + output = torch.divide(total, count).to(dtype=dtype) + if take_sqrt: + output = torch.sqrt(output) + return output + else: + raise ValueError( + f"masked std/var expects strided tensor (got {input.layout} tensor)" + ) + + +@_apply_docstring_templates +def var( + input: Union[Tensor, MaskedTensor], + dim: DimOrDims = None, + unbiased: Optional[bool] = None, + *, + correction: Optional[Union[int, float]] = None, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + """\ +{reduction_signature} +{reduction_descr} +The identity value of sample variance operation is undefined. The +elements of output tensor with strided layout, that correspond to +fully masked-out elements, have ``nan`` values. +{reduction_args} +{reduction_example}""" + return _std_var( + input=input, + dim=dim, + unbiased=unbiased, + correction_opt=correction, + keepdim=keepdim, + dtype=dtype, + mask=mask, + take_sqrt=False, + ) + + +@_apply_docstring_templates +def std( + input: Union[Tensor, MaskedTensor], + dim: DimOrDims = None, + unbiased: Optional[bool] = None, + *, + correction: Optional[int] = None, + keepdim: Optional[bool] = False, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + """\ +{reduction_signature} +{reduction_descr} +The identity value of sample standard deviation operation is undefined. The +elements of output tensor with strided layout, that correspond to +fully masked-out elements, have ``nan`` values. +{reduction_args} +{reduction_example}""" + return _std_var( + input=input, + dim=dim, + unbiased=unbiased, + correction_opt=correction, + keepdim=keepdim, + dtype=dtype, + mask=mask, + take_sqrt=True, + ) + + +@_apply_docstring_templates +def softmax( + input: Union[Tensor, MaskedTensor], + dim: int, + *, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + if dtype is None: + dtype = input.dtype + dim_ = _canonical_dim(dim, input.ndim)[0] + mask_input = _combine_input_and_mask(amax, input, mask) + if mask_input.layout == torch.strided: + return torch.nn.functional.softmax(mask_input, dim_, dtype=dtype) + else: + raise ValueError( + f"masked softmax expects strided tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def log_softmax( + input: Union[Tensor, MaskedTensor], + dim: int, + *, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + if dtype is None: + dtype = input.dtype + dim_ = _canonical_dim(dim, input.ndim)[0] + mask_input = _combine_input_and_mask(amax, input, mask) + if mask_input.layout == torch.strided: + return torch.nn.functional.log_softmax(mask_input, dim_, dtype=dtype) + else: + raise ValueError( + f"masked log_softmax expects strided tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def softmin( + input: Union[Tensor, MaskedTensor], + dim: int, + *, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + if dtype is None: + dtype = input.dtype + dim_ = _canonical_dim(dim, input.ndim)[0] + mask_input = _combine_input_and_mask(amin, input, mask) + if mask_input.layout == torch.strided: + return torch.nn.functional.softmin(mask_input, dim_, dtype=dtype) + else: + raise ValueError( + f"masked softmin expects strided tensor (got {mask_input.layout} tensor)" + ) + + +@_apply_docstring_templates +def normalize( + input: Union[Tensor, MaskedTensor], + ord: float, + dim: int, + *, + eps: float = 1e-12, + dtype: Optional[DType] = None, + mask: Optional[Tensor] = None, +) -> Tensor: + if dtype is None: + dtype = input.dtype + # TODO: eliminate mask_input as unnecessary when using masked divide. + mask_input = _combine_input_and_mask(sum, input, mask) + if mask_input.layout == torch.strided: + nrm_ = norm(input, ord, dim, keepdim=True, dtype=dtype, mask=mask) + # TODO: replace torch.maximum with masked maximum when available. + denom = torch.maximum(nrm_, nrm_.new_full([], eps)) + # TODO: replace torch.divide with masked divide when available. + return torch.divide(mask_input, denom) + else: + raise ValueError( + f"masked normalize expects strided tensor (got {mask_input.layout} tensor)" + ) diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/__init__.py b/lib/python3.12/site-packages/torch/masked/maskedtensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e38e03c87086cf50d031dd5591f64f65399d6ac1 --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/maskedtensor/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates +# flake8: noqa + +from .binary import _apply_native_binary, _is_native_binary +from .core import is_masked_tensor, MaskedTensor +from .passthrough import _apply_pass_through_fn, _is_pass_through_fn +from .reductions import _apply_reduction, _is_reduction +from .unary import _apply_native_unary, _is_native_unary diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d9a82cdbb05d52eae2a6d37da738228823af115 Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/_ops_refs.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/_ops_refs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba8e5ba5d6e0c0d985c106ad4b7d8380d2d71ae6 Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/_ops_refs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/binary.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/binary.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a0086998fc8048e1a9496d0db0414d6d475c4bb Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/binary.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/core.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/core.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dde8dcc1a0cd3864e2a2d7edb12b9f2c3b845ce5 Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/core.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/creation.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/creation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a5dd5d157e727728235749386eaa374f26712d9 Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/creation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/passthrough.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/passthrough.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..045789a6a1b724a809623120d4a74c33d510cd5d Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/passthrough.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/reductions.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/reductions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..43df39dc981dcde352837c4deed974bbfc71bcf7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/reductions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/unary.cpython-312.pyc b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/unary.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4dcd8ba600a95c0c352c1aef67d6856e01e95c81 Binary files /dev/null and b/lib/python3.12/site-packages/torch/masked/maskedtensor/__pycache__/unary.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/_ops_refs.py b/lib/python3.12/site-packages/torch/masked/maskedtensor/_ops_refs.py new file mode 100644 index 0000000000000000000000000000000000000000..d0cb64fa3c7bdc96e9ebbb458a9827f02d424393 --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/maskedtensor/_ops_refs.py @@ -0,0 +1,531 @@ +# mypy: allow-untyped-defs +# Copyright (c) Meta Platforms, Inc. and affiliates + +from functools import partial +from typing import Any, Callable, Dict, TYPE_CHECKING + +import torch + +from .binary import _apply_native_binary, NATIVE_BINARY_FNS, NATIVE_INPLACE_BINARY_FNS +from .core import ( + _get_data, + _masks_match, + _maybe_get_mask, + is_masked_tensor, + MaskedTensor, +) +from .passthrough import _apply_pass_through_fn, PASSTHROUGH_FNS +from .reductions import ( + _apply_reduction, + NATIVE_REDUCE_FNS, + TENSOR_REDUCE_FNS, + TORCH_REDUCE_FNS, +) +from .unary import _apply_native_unary, NATIVE_INPLACE_UNARY_FNS, NATIVE_UNARY_FNS + + +if TYPE_CHECKING: + from torch._ops import OpOverload + + +__all__ = [] # type: ignore[var-annotated] + + +def _check_args_kwargs_length( + args, kwargs, error_prefix, len_args=None, len_kwargs=None +): + if len_args is not None and len_args != len(args): + raise ValueError( + f"{error_prefix}: len(args) must be {len_args} but got {len(args)}" + ) + if len_kwargs is not None and len_kwargs != len(kwargs): + raise ValueError( + f"{error_prefix}: len(kwargs) must be {len_kwargs} but got {len(kwargs)}" + ) + + +class _MaskedContiguous(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + if not is_masked_tensor(input): + raise ValueError("MaskedContiguous forward: input must be a MaskedTensor.") + + if input.is_contiguous(): + return input + + data = input.get_data() + mask = input.get_mask() + + return MaskedTensor(data.contiguous(), mask.contiguous()) + + @staticmethod + def backward(ctx, grad_output): + return grad_output + + +class _MaskedToDense(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + if not is_masked_tensor(input): + raise ValueError("MaskedToDense forward: input must be a MaskedTensor.") + + if input.layout == torch.strided: + return input + + ctx.layout = input.layout + data = input.get_data() + mask = input.get_mask() + + return MaskedTensor(data.to_dense(), mask.to_dense()) + + @staticmethod + def backward(ctx, grad_output): + layout = ctx.layout + + if layout == torch.sparse_coo: + return grad_output.to_sparse_coo() + elif layout == torch.sparse_csr: + return grad_output.to_sparse_csr() + elif layout == torch.strided: + return grad_output.to_dense() + raise ValueError("to_dense: Unsupported input layout: ", layout) + + +class _MaskedToSparse(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + if not is_masked_tensor(input): + raise ValueError("MaskedToSparse forward: input must be a MaskedTensor.") + + # Following the convention from sparse tensors that to_sparse always means that we convert to sparse_coo + if input.layout == torch.sparse_coo: + return input + + data = input.get_data() + mask = input.get_mask() + sparse_mask = mask.to_sparse_coo().coalesce() + sparse_data = data.sparse_mask(sparse_mask) + + return MaskedTensor(sparse_data, sparse_mask) + + @staticmethod + def backward(ctx, grad_output): + return grad_output.to_dense() + + +class _MaskedToSparseCsr(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + if not is_masked_tensor(input): + raise ValueError("MaskedToSparseCsr forward: input must be a MaskedTensor.") + + if input._masked_data.ndim != 2: + raise ValueError( + f"Only 2D tensors can be converted to the SparseCsr layout but got shape: {input._masked_data.size()}" + ) + + if input.layout == torch.sparse_csr: + return input + + data = input.get_data() + mask = input.get_mask() + sparse_mask = mask.to_sparse_csr() + sparse_data = data.sparse_mask(sparse_mask) + + return MaskedTensor(sparse_data, sparse_mask) + + @staticmethod + def backward(ctx, grad_output): + return grad_output.to_dense() + + +class _MaskedWhere(torch.autograd.Function): + @staticmethod + def forward(ctx, cond, self, other): + ctx.mark_non_differentiable(cond) + ctx.save_for_backward(cond) + return torch.ops.aten.where(cond, self, other) + + @staticmethod + def backward(ctx, grad_output): + (cond,) = ctx.saved_tensors + + def masked_out_like(mt): + return MaskedTensor(mt.get_data(), torch.zeros_like(mt.get_mask()).bool()) + + return ( + None, + torch.ops.aten.where(cond, grad_output, masked_out_like(grad_output)), + torch.ops.aten.where(cond, masked_out_like(grad_output), grad_output), + ) + + +_MASKEDTENSOR_FUNCTION_TABLE = {} + +_function_fn_apply_map = { + ( + tuple(NATIVE_REDUCE_FNS), + tuple(TORCH_REDUCE_FNS), + tuple(TENSOR_REDUCE_FNS), + ): _apply_reduction, +} + +for fn_map_list, apply_fn in _function_fn_apply_map.items(): + for fn_map in fn_map_list: + for fn in fn_map: + _MASKEDTENSOR_FUNCTION_TABLE[fn] = partial(apply_fn, fn) + + +def register_function_func(ops): + """ + Used for registering a new __torch_function__ function to MaskedTensor + Called via _MASKEDTENSOR_FUNCTION_TABLE[func](*args, **kwargs) + + The code to register a new function looks like: + + @register_function_func(list_of_ops) + def foo(func, *args, **kwargs): + + """ + + def wrapper(func): + for op in ops: + _MASKEDTENSOR_FUNCTION_TABLE[op] = partial(func, op) + + return wrapper + + +@register_function_func(NATIVE_REDUCE_FNS + TORCH_REDUCE_FNS + TENSOR_REDUCE_FNS) +def _general_function_reductions(func, *args, **kwargs): + return _apply_reduction(func, *args, **kwargs) + + +@register_function_func([torch.Tensor.where, torch.where]) +def _function_where(func, *args, **kwargs): + _check_args_kwargs_length( + args, kwargs, "__torch_function__, torch.where", len_args=3, len_kwargs=0 + ) + return _MaskedWhere.apply(*args) + + +@register_function_func([torch.Tensor.contiguous]) +def _function_contiguous(func, *args, **kwargs): + return _MaskedContiguous.apply(args[0]) + + +@register_function_func([torch.Tensor.to_dense]) +def _function_to_dense(func, *args, **kwargs): + return _MaskedToDense.apply(args[0]) + + +@register_function_func([torch.Tensor.to_sparse]) +def _function_to_sparse(func, *args, **kwargs): + return _MaskedToSparse.apply(args[0]) + + +@register_function_func([torch.Tensor.to_sparse_csr]) +def _function_to_sparse_csr(func, *args, **kwargs): + return _MaskedToSparseCsr.apply(args[0]) + + +_MASKEDTENSOR_DISPATCH_TABLE: Dict["OpOverload", Callable[..., Any]] = {} + + +def register_dispatch_func(aten_ops): + """ + Used for registering a new __torch_dispatch__ function to MaskedTensor + Called via _MASKEDTENSOR_DISPATCH_TABLE[func](*args, **kwargs) + + The code to register a new function looks like: + + @register_dispatch_func(list_of_ops) + def foo(func, *args, **kwargs): + + """ + + def wrapper(func): + for aten_op in aten_ops: + _MASKEDTENSOR_DISPATCH_TABLE[aten_op] = partial(func, aten_op) + + return wrapper + + +@register_dispatch_func(NATIVE_REDUCE_FNS + TORCH_REDUCE_FNS + TENSOR_REDUCE_FNS) +def _general_reduction(func, *args, **kwargs): + return _apply_reduction(func, *args, **kwargs) + + +@register_dispatch_func(PASSTHROUGH_FNS) +def _general_passthrough(func, *args, **kwargs): + return _apply_pass_through_fn(func, *args, **kwargs) + + +@register_dispatch_func(NATIVE_UNARY_FNS + NATIVE_INPLACE_UNARY_FNS) +def _general_unary(func, *args, **kwargs): + return _apply_native_unary(func, *args, **kwargs) + + +@register_dispatch_func(NATIVE_BINARY_FNS + NATIVE_INPLACE_BINARY_FNS) +def _general_binary(func, *args, **kwargs): + return _apply_native_binary(func, *args, **kwargs) + + +@register_dispatch_func([torch.ops.aten.stride]) +def stride(func, *args, **kwargs): + return None + + +@register_dispatch_func([torch.ops.aten.sym_stride]) +def sym_stride(func, *args, **kwargs): + return None + + +@register_dispatch_func([torch.ops.prim.layout]) +def layout(func, *args, **kwargs): + return _get_data(args[0]).layout + + +@register_dispatch_func([torch.ops.aten.is_contiguous]) +def is_contiguous(func, *args, **kwargs): + data = _get_data(args[0]) + if data.is_sparse: + raise ValueError("MaskedTensors with sparse data do not have is_contiguous") + return func(data, *args[1:], **kwargs) + + +@register_dispatch_func([torch.ops.aten.is_strides_like_format]) +def is_strides_like_format(func, *args, **kwargs): + data = _get_data(args[0]) + if data.is_sparse: + raise ValueError( + "MaskedTensors with sparse data do not have is_strides_like_format" + ) + return func(data, *args[1:], **kwargs) + + +@register_dispatch_func([torch.ops.aten.is_non_overlapping_and_dense]) +def is_non_overlapping_and_dense(func, *args, **kwargs): + data = _get_data(args[0]) + if data.is_sparse: + raise ValueError( + "MaskedTensors with sparse data do not have is_non_overlapping_and_dense" + ) + return func(data, *args[1:], **kwargs) + + +@register_dispatch_func([torch.ops.aten.contiguous]) +def contiguous(func, *args, **kwargs): + if _get_data(args[0]).is_sparse: + raise ValueError("MaskedTensors with sparse data do not have contiguous") + return _MaskedContiguous.apply(args[0]) + + +@register_dispatch_func([torch.ops.aten.new_empty_strided]) +def new_empty_strided(func, *args, **kwargs): + _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=3) + data = _get_data(args[0]) + mask = _maybe_get_mask(args[0]) + if tuple(args[1]) != tuple(data.size()): + raise ValueError( + f"__torch_dispatch__, {func}: args[1] expected to be the same as data.size()" + ) + if tuple(args[2]) != tuple(data.stride()): + raise ValueError( + f"__torch_dispatch__, {func}: args[2] expected to be the same as data.stride()" + ) + return MaskedTensor(func(data, args[1], args[2], **kwargs), mask) + + +@register_dispatch_func([torch.ops.aten._local_scalar_dense]) +def _local_scalar_dense(func, *args, **kwargs): + if not _maybe_get_mask(args[0]): + raise ValueError(f"__torch_dispatch__, {func}: expected a mask tensor") + return torch.ops.aten._local_scalar_dense(_get_data(args[0])) + + +@register_dispatch_func([torch.ops.aten.detach, torch.ops.aten.clone]) +def _apply_fn_on_data(func, *args, **kwargs): + return MaskedTensor(func(_get_data(args[0])), _maybe_get_mask(args[0])) + + +@register_dispatch_func([torch.ops.aten._to_copy]) +def _to_copy(func, *args, **kwargs): + new_data = func(_get_data(args[0]), *args[1:], **kwargs) + return MaskedTensor(new_data, _maybe_get_mask(args[0])) + + +@register_dispatch_func([torch.ops.aten._softmax]) +def _softmax(func, *args, **kwargs): + _check_args_kwargs_length( + args, kwargs, f"__torch_dispatch__, {func}", len_args=3, len_kwargs=0 + ) + data = _get_data(args[0]) + mask = _maybe_get_mask(args[0]) + result_data = torch.ops.aten._masked_softmax(data, ~mask, args[1], 2) + return MaskedTensor(result_data, mask) + + +@register_dispatch_func([torch.ops.aten.ones_like]) +def ones_like(func, *args, **kwargs): + _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=1) + result_data = func(_get_data(args[0]), **kwargs) + return MaskedTensor(result_data, _maybe_get_mask(args[0])) + + +@register_dispatch_func([torch.ops.aten._softmax_backward_data]) +def _softmax_backward_data(func, *args, **kwargs): + _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=4) + grad, output, dim, _input_dtype = args + if is_masked_tensor(grad) and is_masked_tensor(output): + if not _masks_match(grad, output): + raise ValueError( + "__torch_dispatch__, {func}: expected the masks of grad and output to match" + ) + grad_data = _get_data(grad) + new_grad_data = torch.ops.aten._masked_softmax_backward( + grad_data, + _get_data(output), + ~_maybe_get_mask(grad), + dim % grad_data.ndim, + ) + res = MaskedTensor(new_grad_data, _maybe_get_mask(grad)) + return res + else: + raise ValueError( + f"__torch_dispatch__, {func}: grad and output must both be MaskedTensors" + ) + + +@register_dispatch_func([torch.ops.aten.copy_]) +def copy_(func, *args, **kwargs): + _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=2) + if not _masks_match(_maybe_get_mask(args[0]), _maybe_get_mask(args[1])): + raise ValueError("args[0] mask and args[1] mask must match but do not") + func(_get_data(args[0]), _get_data(args[1])) + return args[0] + + +@register_dispatch_func([torch.ops.aten.where]) +def where(func, *args, **kwargs): + _check_args_kwargs_length( + args, kwargs, f"__torch_dispatch__, {func}", len_args=3, len_kwargs=0 + ) + if not torch.is_tensor(args[0]): + raise ValueError("__torch_dispatch__, {func}: expected args[0] to be a tensor") + mx = args[1] + my = args[2] + if not is_masked_tensor(mx): + mx = MaskedTensor(mx, torch.ones_like(mx, dtype=torch.bool)) + if not is_masked_tensor(my): + my = MaskedTensor(my, torch.ones_like(my, dtype=torch.bool)) + new_data = func(args[0], mx.get_data(), my.get_data()) + new_mask = func(args[0], mx.get_mask(), my.get_mask()) + return MaskedTensor(new_data, new_mask) + + +@register_dispatch_func([torch.ops.aten._to_sparse]) +def _to_sparse(func, *args, **kwargs): + _check_args_kwargs_length( + args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0 + ) + if not torch.is_tensor(args[0]): + raise TypeError("__torch_dispatch__, {func}: expected args[0] to be a tensor") + mt = args[0] + if not is_masked_tensor(mt): + mt = MaskedTensor(mt, torch.ones_like(mt, dtype=torch.bool)) + if mt.is_sparse_coo(): + return mt + new_mask = func(_maybe_get_mask(args[0])).coalesce() + new_data = _get_data(args[0]).sparse_mask(new_mask) + return MaskedTensor(new_data, new_mask) + + +@register_dispatch_func([torch.ops.aten._to_sparse_csr]) +def _to_sparse_csr(func, *args, **kwargs): + _check_args_kwargs_length( + args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0 + ) + if not torch.is_tensor(args[0]): + raise ValueError("__torch_dispatch__, {func}: expected args[0] to be a tensor") + mt = args[0] + if not is_masked_tensor(mt): + mt = MaskedTensor(mt, torch.ones_like(mt).bool()) + if mt.is_sparse_csr(): + return mt + new_mask = func(_maybe_get_mask(args[0])) + new_data = _get_data(args[0]).sparse_mask(new_mask) + return MaskedTensor(new_data, new_mask) + + +@register_dispatch_func([torch.ops.aten._to_dense]) +def _to_dense(func, *args, **kwargs): + _check_args_kwargs_length( + args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0 + ) + if not torch.is_tensor(args[0]): + raise ValueError("__torch_dispatch__, {func}: expected args[0] to be a tensor") + mt = args[0] + if not is_masked_tensor(mt): + mt = MaskedTensor(mt, torch.ones_like(mt).bool()) + new_data = func(_get_data(args[0])) + new_mask = func(_maybe_get_mask(args[0])) + return MaskedTensor(new_data, new_mask) + + +@register_dispatch_func([torch.ops.aten._indices]) +def _indices(func, *args, **kwargs): + # Assumes data is sparse + _check_args_kwargs_length( + args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0 + ) + data = _get_data(args[0]).indices() + return MaskedTensor(data, torch.ones_like(data).bool()) + + +@register_dispatch_func([torch.ops.aten._values]) +def _values(func, *args, **kwargs): + _check_args_kwargs_length( + args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0 + ) + data = _get_data(args[0]).values() + return MaskedTensor(data, torch.ones_like(data).bool()) + + +@register_dispatch_func([torch.ops.aten._sparse_coo_tensor_with_dims_and_tensors]) +def _sparse_coo_tensor_with_dims_and_tensors(func, *args, **kwargs): + new_args = list(args) + if is_masked_tensor(args[-1]): + new_args[-1] = args[-1].get_data() + if is_masked_tensor(args[-2]): + new_args[-2] = args[-2].get_data() + + new_data = func(*new_args, **kwargs) + new_args[-1] = torch.ones_like(new_args[-1]) + new_mask = func(*new_args, **kwargs).bool() + + return MaskedTensor(new_data, new_mask) + + +@register_dispatch_func([torch.ops.aten.is_same_size]) +def is_same_size(func, *args, **kwargs): + _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=2) + return _get_data(args[0]).is_same_size(_get_data(args[1])) + + +@register_dispatch_func([torch.ops.aten._is_any_true]) +def _is_any_true(func, *args, **kwargs): + _check_args_kwargs_length( + args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0 + ) + data = _get_data(args[0]) + mask = _maybe_get_mask(args[0]) + if mask is None: + raise ValueError( + f"__torch_dispatch__, {func}: expected args[0] to be a MaskedTensor" + ) + if data.dtype != torch.bool: + raise ValueError(f"__torch_dispatch__, {func}: expected a boolean tensor") + if data.is_sparse: + raise ValueError(f"MaskedTensors with sparse data do not have {func}") + + return MaskedTensor(func(data & mask), torch.tensor(True)) diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/binary.py b/lib/python3.12/site-packages/torch/masked/maskedtensor/binary.py new file mode 100644 index 0000000000000000000000000000000000000000..8315ae11be7175c2b5aaef178a4bc4785dcbcb29 --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/maskedtensor/binary.py @@ -0,0 +1,200 @@ +# mypy: allow-untyped-defs +# Copyright (c) Meta Platforms, Inc. and affiliates + +import torch + +from .core import ( + _map_mt_args_kwargs, + _masks_match, + _tensors_match, + _wrap_result, + is_masked_tensor, +) + + +__all__ = [] # type: ignore[var-annotated] + +BINARY_NAMES = [ + "add", + "atan2", + "arctan2", + "bitwise_and", + "bitwise_or", + "bitwise_xor", + "bitwise_left_shift", + "bitwise_right_shift", + "div", + "divide", + "floor_divide", + "fmod", + "logaddexp", + "logaddexp2", + "mul", + "multiply", + "nextafter", + "remainder", + "sub", + "subtract", + "true_divide", + "eq", + "ne", + "le", + "ge", + "greater", + "greater_equal", + "gt", + "less_equal", + "lt", + "less", + "maximum", + "minimum", + "fmax", + "fmin", + "not_equal", +] + +INPLACE_BINARY_NAMES = [ + n + "_" + for n in ( + list( + set(BINARY_NAMES) + - { + "logaddexp", + "logaddexp2", + "equal", + "fmin", + "minimum", + "maximum", + "fmax", + } + ) + ) +] + + +def _get_at_least_one_mask(a, b): + if not is_masked_tensor(a) and not is_masked_tensor(b): + raise TypeError("At least one of `a` and `b` must be a MaskedTensor") + if not _masks_match(a, b): + raise ValueError("a and b must have matching masks") + if is_masked_tensor(a): + return a.get_mask() + return b.get_mask() + + +def _binary_helper(fn, args, kwargs, inplace): + if len(kwargs) != 0: + raise ValueError("len(kwargs) must equal 0") + for a in args[2:]: + if torch.is_tensor(a): + raise TypeError( + "MaskedTensor binary ops do not support Tensor arguments aside from the lhs and rhs" + ) + + if not _masks_match(*args[:2]): + raise ValueError( + "Input masks must match. If you need support for this, please open an issue on Github." + ) + + data_args, _data_kwargs = _map_mt_args_kwargs(args, kwargs, lambda x: x.get_data()) + mask_args, _mask_kwargs = _map_mt_args_kwargs(args, kwargs, lambda x: x.get_mask()) + + args0_layout = data_args[0].layout + same_layout = ( + torch.is_tensor(data_args[1]) or is_masked_tensor(data_args[1]) + ) and (args0_layout == data_args[1].layout) + + if args0_layout == torch.sparse_coo: + if same_layout: + if not _tensors_match(data_args[0].indices(), data_args[1].indices()): + raise ValueError( + "sparse_coo indices must match. If you need support for this, please open an issue on Github." + ) + if data_args[0].size() != data_args[1].size(): + raise ValueError( + "input1 and input2 must have the same size for binary functions." + ) + + data_args[1] = data_args[1].values() + + i = data_args[0].indices() + size = data_args[0].size() + data_args[0] = data_args[0].values() + v = fn(*data_args) + result_data = torch.sparse_coo_tensor(i, v, size) + + elif args0_layout == torch.sparse_csr: + if same_layout: + if not ( + _tensors_match(data_args[0].crow_indices(), data_args[1].crow_indices()) + and _tensors_match( + data_args[0].col_indices(), data_args[1].col_indices() + ) + ): + raise ValueError( + "sparse_csr indices must match. If you need support for this, please open an issue on Github." + ) + + data_args[1] = data_args[1].values() + + crow = data_args[0].crow_indices() + col = data_args[0].col_indices() + size = data_args[0].size() + data_args[0] = data_args[0].values() + v = fn(*data_args) + result_data = torch.sparse_csr_tensor(crow, col, v, size) + + else: + result_data = fn(*data_args) + + if inplace: + args[0]._set_data_mask(result_data, mask_args[0]) + return args[0] + else: + result_mask = _get_at_least_one_mask(*args[:2]) + # sparse tensors don't have strides so we can only expand if the layout is strided + if args0_layout == torch.strided: + result_mask = result_mask.expand_as(result_data) + return _wrap_result(result_data, result_mask) + + +def _torch_binary(fn_name): + fn = getattr(torch.ops.aten, fn_name) + + def binary_fn(*args, **kwargs): + return _binary_helper(fn, args, kwargs, inplace=False) + + return binary_fn + + +def _torch_inplace_binary(fn_name): + fn = getattr(torch.ops.aten, fn_name) + + def binary_fn(*args, **kwargs): + return _binary_helper(fn, args, kwargs, inplace=True) + + return binary_fn + + +NATIVE_BINARY_MAP = { + getattr(torch.ops.aten, name): _torch_binary(name) for name in BINARY_NAMES +} +NATIVE_INPLACE_BINARY_MAP = { + getattr(torch.ops.aten, name): _torch_inplace_binary(name) + for name in INPLACE_BINARY_NAMES +} + +NATIVE_BINARY_FNS = list(NATIVE_BINARY_MAP.keys()) +NATIVE_INPLACE_BINARY_FNS = list(NATIVE_INPLACE_BINARY_MAP.keys()) + + +def _is_native_binary(fn): + return fn in NATIVE_BINARY_FNS or fn in NATIVE_INPLACE_BINARY_FNS + + +def _apply_native_binary(fn, *args, **kwargs): + if fn in NATIVE_BINARY_FNS: + return NATIVE_BINARY_MAP[fn](*args, **kwargs) + if fn in NATIVE_INPLACE_BINARY_FNS: + return NATIVE_INPLACE_BINARY_MAP[fn](*args, **kwargs) + return NotImplemented diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/core.py b/lib/python3.12/site-packages/torch/masked/maskedtensor/core.py new file mode 100644 index 0000000000000000000000000000000000000000..d1cc620325933375b7020a2918aa5b43bb9a8377 --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/maskedtensor/core.py @@ -0,0 +1,359 @@ +# mypy: allow-untyped-defs +# Copyright (c) Meta Platforms, Inc. and affiliates + +import warnings +from typing import Any +from typing_extensions import TypeIs + +import torch +from torch.overrides import get_default_nowrap_functions + + +__all__ = [ + "MaskedTensor", + "is_masked_tensor", +] + + +def is_masked_tensor(obj: Any, /) -> TypeIs["MaskedTensor"]: + r"""Returns True if the input is a MaskedTensor, else False + + Args: + a: any input + + Examples: + + >>> # xdoctest: +SKIP + >>> from torch.masked import MaskedTensor + >>> data = torch.arange(6).reshape(2,3) + >>> mask = torch.tensor([[True, False, False], [True, True, False]]) + >>> mt = MaskedTensor(data, mask) + >>> is_masked_tensor(mt) + True + """ + return isinstance(obj, MaskedTensor) + + +def _tensors_match(a, b, exact=True, rtol=1e-05, atol=1e-08): + if is_masked_tensor(a) or is_masked_tensor(b): + raise ValueError("Neither `a` nor `b` can be a MaskedTensor.") + if a.layout != b.layout: + raise ValueError( + f"`a` and `b` must have the same layout. Got {a.layout} and {b.layout}" + ) + + if a.dtype != b.dtype: + b = b.type(a.dtype) + if a.layout == b.layout == torch.sparse_coo: + return _tensors_match(a.values(), b.values(), exact) and _tensors_match( + a.indices(), b.indices(), exact + ) + elif a.layout == b.layout == torch.sparse_csr: + return ( + _tensors_match(a.crow_indices(), b.crow_indices(), exact) + and _tensors_match(a.col_indices(), b.col_indices(), exact) + and _tensors_match(a.values(), b.values(), exact) + ) + if exact: + return (a.dim() == b.dim()) and torch.eq(a, b).all().item() + return (a.dim() == b.dim()) and torch.allclose(a, b, rtol=rtol, atol=atol) + + +def _masks_match(a, b): + if is_masked_tensor(a) and is_masked_tensor(b): + mask_a = a.get_mask() + mask_b = b.get_mask() + return _tensors_match(mask_a, mask_b, exact=True) + return True + + +def _map_mt_args_kwargs(args, kwargs, map_fn): + def _helper(a, map_fn): + if is_masked_tensor(a): + return map_fn(a) + elif torch.is_tensor(a): + return a + elif isinstance(a, list): + a_impl, _ = _map_mt_args_kwargs(a, {}, map_fn) + return a_impl + elif isinstance(a, tuple): + a_impl, _ = _map_mt_args_kwargs(a, {}, map_fn) + return tuple(a_impl) + else: + return a + + if kwargs is None: + kwargs = {} + impl_args = [] + for a in args: + impl_args.append(_helper(a, map_fn)) + impl_kwargs = {} + for k in kwargs.keys(): + impl_kwargs[k] = _helper(a, map_fn) + return impl_args, impl_kwargs + + +def _wrap_result(result_data, result_mask): + if isinstance(result_data, list): + return [_wrap_result(r, m) for (r, m) in zip(result_data, result_mask)] + if isinstance(result_data, tuple): + return tuple(_wrap_result(r, m) for (r, m) in zip(result_data, result_mask)) + if torch.is_tensor(result_data): + return MaskedTensor(result_data, result_mask) + # Expect result_data and result_mask to be Tensors only + return NotImplemented + + +def _masked_tensor_str(data, mask, formatter): + if data.layout in {torch.sparse_coo, torch.sparse_csr}: + data = data.to_dense() + mask = mask.to_dense() + if data.dim() == 1: + formatted_elements = [ + formatter.format(d.item()) if isinstance(d.item(), float) else str(d.item()) + for d in data + ] + max_len = max(8 if x[1] else len(x[0]) for x in zip(formatted_elements, ~mask)) + return ( + "[" + + ", ".join( + [ + "--".rjust(max_len) if m else e + for (e, m) in zip(formatted_elements, ~mask) + ] + ) + + "]" + ) + sub_strings = [_masked_tensor_str(d, m, formatter) for (d, m) in zip(data, mask)] + sub_strings = ["\n".join([" " + si for si in s.split("\n")]) for s in sub_strings] + return "[\n" + ",\n".join(sub_strings) + "\n]" + + +def _get_data(a): + if is_masked_tensor(a): + return a._masked_data + return a + + +def _maybe_get_mask(a): + if is_masked_tensor(a): + return a.get_mask() + return None + + +class MaskedTensor(torch.Tensor): + @staticmethod + def __new__(cls, data, mask, requires_grad=False): + if is_masked_tensor(data) or not torch.is_tensor(data): + raise TypeError("data must be a Tensor") + if is_masked_tensor(mask) or not torch.is_tensor(mask): + raise TypeError("mask must be a Tensor") + # Use a Tensor that of the give size for the wrapper. + kwargs = { + "device": data.device, + "dtype": data.dtype, + "layout": data.layout, + "requires_grad": requires_grad, + "dispatch_sizes_strides_policy": "strides", + "dispatch_layout": True, + } + warnings.warn( + ( + "The PyTorch API of MaskedTensors is in prototype stage " + "and will change in the near future. Please open a Github issue " + "for features requests and see our documentation on the torch.masked " + "module for further information about the project." + ), + UserWarning, + stacklevel=2, + ) + if data.requires_grad: + warnings.warn( + "It is not recommended to create a MaskedTensor with a tensor that requires_grad. " + "To avoid this, you can use data.clone().detach()", + UserWarning, + stacklevel=2, + ) + return torch.Tensor._make_wrapper_subclass(cls, data.size(), **kwargs) # type: ignore[attr-defined] + + def _preprocess_data(self, data, mask): + from .._ops import _sparse_coo_where, _sparse_csr_where + + if data.layout != mask.layout: + raise TypeError("data and mask must have the same layout.") + if data.layout == torch.sparse_coo: + data = data.coalesce() + mask = mask.coalesce() + if data._nnz() != mask._nnz(): + data = _sparse_coo_where(mask, data, torch.tensor(0)) + elif data.layout == torch.sparse_csr: + if data._nnz() != mask._nnz(): + data = _sparse_csr_where(mask, data, torch.tensor(0)) + + # Have to pick awkward names to not conflict with existing fields such as data + self._masked_data = data.clone() + self._masked_mask = mask.clone() + + def _validate_members(self): + data = self._masked_data + mask = self.get_mask() + if type(data) != type(mask): + raise TypeError( + f"data and mask must have the same type. Got {type(data)} and {type(mask)}" + ) + if data.layout not in {torch.strided, torch.sparse_coo, torch.sparse_csr}: + raise TypeError(f"data layout of {data.layout} is not supported.") + if data.layout == torch.sparse_coo: + if not _tensors_match(data.indices(), mask.indices(), exact=True): + raise ValueError( + "data and mask are both sparse COO tensors but do not have the same indices." + ) + elif data.layout == torch.sparse_csr: + if not _tensors_match( + data.crow_indices(), mask.crow_indices(), exact=True + ) or not _tensors_match(data.col_indices(), mask.col_indices(), exact=True): + raise ValueError( + "data and mask are both sparse CSR tensors but do not share either crow or col indices." + ) + if mask.dtype != torch.bool: + raise TypeError("mask must have dtype bool.") + if not ( + data.dtype == torch.float16 + or data.dtype == torch.float32 + or data.dtype == torch.float64 + or data.dtype == torch.bool + or data.dtype == torch.int8 + or data.dtype == torch.int16 + or data.dtype == torch.int32 + or data.dtype == torch.int64 + ): + raise TypeError(f"{data.dtype} is not supported in MaskedTensor.") + if data.dim() != mask.dim(): + raise ValueError("data.dim() must equal mask.dim()") + if data.size() != mask.size(): + raise ValueError("data.size() must equal mask.size()") + + def __init__(self, data, mask, requires_grad=False): + self._preprocess_data(data, mask) + self._validate_members() + + @staticmethod + def _from_values(data, mask): + """Differentiable constructor for MaskedTensor""" + + class Constructor(torch.autograd.Function): + @staticmethod + def forward(ctx, data, mask): + return MaskedTensor(data, mask) + + @staticmethod + def backward(ctx, grad_output): + return grad_output, None + + result = Constructor.apply(data, mask) + return result + + def _set_data_mask(self, data, mask): + self._masked_data = data + self._masked_mask = mask + self._validate_members() + + def __repr__(self): # type: ignore[override] + formatter = "{0:8.4f}" + if self.dim() == 0: + scalar_data = self.get_data().item() + data_formatted = ( + formatter.format(scalar_data) + if isinstance(scalar_data, float) + else str(scalar_data) + ) + if not self.get_mask().item(): + data_formatted = "--" + return ( + "MaskedTensor(" + + data_formatted + + ", " + + str(self.get_mask().item()) + + ")" + ) + s = _masked_tensor_str(self.get_data(), self.get_mask(), formatter) + s = "\n".join(" " + si for si in s.split("\n")) + return "MaskedTensor(\n" + s + "\n)" + + # Seems like this needs to be defined before torch_dispatch to work + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + + from ._ops_refs import _MASKEDTENSOR_FUNCTION_TABLE + + if func in _MASKEDTENSOR_FUNCTION_TABLE: + return _MASKEDTENSOR_FUNCTION_TABLE[func](*args, **kwargs) + + if not all(issubclass(cls, t) for t in types): + return NotImplemented + with torch._C.DisableTorchFunctionSubclass(): + ret = func(*args, **kwargs) + if func in get_default_nowrap_functions(): + return ret + else: + return torch._tensor._convert(ret, cls) + + @classmethod + def unary(cls, fn, data, mask): + return MaskedTensor(fn(data), mask) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + func = func.overloadpacket + + from ._ops_refs import _MASKEDTENSOR_DISPATCH_TABLE + + if func in _MASKEDTENSOR_DISPATCH_TABLE: + return _MASKEDTENSOR_DISPATCH_TABLE[func](*args, **kwargs) + + msg = ( + f"{func.__name__} is not implemented in __torch_dispatch__ for MaskedTensor.\n" + "If you would like this operator to be supported, please file an issue for a feature request at " + "https://github.com/pytorch/maskedtensor/issues with a minimal reproducible code snippet.\n" + "In the case that the semantics for the operator are not trivial, it would be appreciated " + "to also include a proposal for the semantics." + ) + warnings.warn(msg) + return NotImplemented + + def __lt__(self, other): + if is_masked_tensor(other): + return MaskedTensor(self.get_data() < _get_data(other), self.get_mask()) + return MaskedTensor(self.get_data() < other, self.get_mask()) + + def to_tensor(self, value): + return self.get_data().masked_fill(~self.get_mask(), value) + + def get_data(self): + class GetData(torch.autograd.Function): + @staticmethod + def forward(ctx, self): + return self._masked_data.detach() + + @staticmethod + def backward(ctx, grad_output): + if is_masked_tensor(grad_output): + return grad_output + return MaskedTensor(grad_output, self.get_mask()) + + return GetData.apply(self) + + def get_mask(self): + return self._masked_mask + + def is_sparse_coo(self): + return self.layout == torch.sparse_coo + + def is_sparse_csr(self): # type: ignore[override] + return self.layout == torch.sparse_csr + + # Update later to support more sparse layouts + @property + def is_sparse(self): + return self.is_sparse_coo() or self.is_sparse_csr() diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/creation.py b/lib/python3.12/site-packages/torch/masked/maskedtensor/creation.py new file mode 100644 index 0000000000000000000000000000000000000000..a013ef1beb6668c5fc2456c2c5a17f2399fc6be8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/maskedtensor/creation.py @@ -0,0 +1,23 @@ +# mypy: allow-untyped-defs +# Copyright (c) Meta Platforms, Inc. and affiliates + +from .core import MaskedTensor + + +__all__ = [ + "as_masked_tensor", + "masked_tensor", +] + + +# These two factory functions are intended to mirror +# torch.tensor - guaranteed to be a leaf node +# torch.as_tensor - differentiable constructor that preserves the autograd history + + +def masked_tensor(data, mask, requires_grad=False): + return MaskedTensor(data, mask, requires_grad) + + +def as_masked_tensor(data, mask): + return MaskedTensor._from_values(data, mask) diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/passthrough.py b/lib/python3.12/site-packages/torch/masked/maskedtensor/passthrough.py new file mode 100644 index 0000000000000000000000000000000000000000..ba13f50c1fee9c9fc10563ffc9f4ff3211c0dca6 --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/maskedtensor/passthrough.py @@ -0,0 +1,50 @@ +# mypy: allow-untyped-defs +# Copyright (c) Meta Platforms, Inc. and affiliates +""" +These are functions that should simply be applied to both mask and data. +Take select or stack as an example. This operation can be applied to +both the mask and data of a MaskedTensor and the result wrapped into +a new MaskedTensor as a result. +""" + +import torch + +from .core import _map_mt_args_kwargs, _wrap_result + + +__all__ = [] # type: ignore[var-annotated] + + +PASSTHROUGH_FNS = [ + torch.ops.aten.select, + torch.ops.aten.transpose, + torch.ops.aten.split, + torch.ops.aten.t, + torch.ops.aten.slice, + torch.ops.aten.slice_backward, + torch.ops.aten.select_backward, + torch.ops.aten.index, + torch.ops.aten.expand, + torch.ops.aten.view, + torch.ops.aten._unsafe_view, + torch.ops.aten._reshape_alias, + torch.ops.aten.cat, + torch.ops.aten.unsqueeze, + torch.ops.aten.unfold, + torch.ops.aten.unfold_backward, + torch.ops.aten.im2col, + torch.ops.aten.col2im, + torch.ops.aten.stack, +] + + +def _is_pass_through_fn(fn): + return fn in PASSTHROUGH_FNS + + +def _apply_pass_through_fn(fn, *args, **kwargs): + data_args, data_kwargs = _map_mt_args_kwargs(args, kwargs, lambda x: x.get_data()) + result_data = fn(*data_args, **data_kwargs) + mask_args, mask_kwargs = _map_mt_args_kwargs(args, kwargs, lambda x: x.get_mask()) + result_mask = fn(*mask_args, **mask_kwargs) + return _wrap_result(result_data, result_mask) diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/reductions.py b/lib/python3.12/site-packages/torch/masked/maskedtensor/reductions.py new file mode 100644 index 0000000000000000000000000000000000000000..fedab1c12a63786fe950edda37ced9637beed83c --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/maskedtensor/reductions.py @@ -0,0 +1,176 @@ +# mypy: allow-untyped-defs +# Copyright (c) Meta Platforms, Inc. and affiliates + +import warnings + +import torch + +from .core import is_masked_tensor +from .creation import as_masked_tensor, masked_tensor + + +__all__ = [] # type: ignore[var-annotated] + + +def _masked_all_all(data, mask=None): + if mask is None: + return data.all() + return data.masked_fill(~mask, True).all() + + +def _masked_all_dim(data, dim, keepdim=False, mask=None): + if mask is None: + return torch.all(data, dim=dim, keepdim=keepdim) + return torch.all(data.masked_fill(~mask, True), dim=dim, keepdim=keepdim) + + +def _masked_all(*args, **kwargs): + if len(args) == 1 and len(kwargs) == 1: + return _masked_all_all(args[0], mask=kwargs["mask"]) + return _masked_all_dim(*args, **kwargs) + + +def _multidim_any(mask, dim, keepdim): + if isinstance(dim, int): + return _multidim_any(mask, [dim], keepdim) + for d in sorted(dim, reverse=True): + mask = torch.any(mask, dim=d, keepdim=keepdim) + return mask + + +def _get_masked_fn(fn): + if fn == "all": + return _masked_all + return getattr(torch.masked, fn) + + +def _torch_reduce_all(fn): + def reduce_all(self): + masked_fn = _get_masked_fn(fn) + data = self.get_data() + mask = self.get_mask().values() if self.is_sparse else self.get_mask() + # When reduction is "all", then torch.argmin/torch.argmax needs to return the index of the + # element corresponding to the min/max, but this operation isn't supported correctly for sparse layouts. + # Therefore, this implementation calculates it using the strides. + if fn == "all": + result_data = masked_fn(data, mask=mask) + + elif fn in {"argmin", "argmax"} and self.is_sparse_coo(): + sparse_idx = masked_fn(data.values(), mask=mask).to(dtype=torch.int) + indices = ( + data.to_sparse_coo().indices() + if not self.is_sparse_coo() + else data.indices() + ) + idx = indices.unbind(1)[sparse_idx] + stride = data.size().numel() / torch.tensor( + data.size(), device=data.device + ).cumprod(0) + result_data = torch.sum(idx * stride) + + # we simply pass in the values for sparse COO/CSR tensors + elif self.is_sparse: + result_data = masked_fn(masked_tensor(data.values(), mask)) + + else: + result_data = masked_fn(self, mask=mask) + + return as_masked_tensor(result_data, torch.any(mask)) + + return reduce_all + + +def _torch_reduce_dim(fn): + def reduce_dim(self, dim, keepdim=False, dtype=None): + if self.is_sparse: + msg = ( + f"The sparse version of {fn} is not implemented in reductions.\n" + "If you would like this operator to be supported, please file an issue for a feature request at " + "https://github.com/pytorch/maskedtensor/issues with a minimal reproducible code snippet.\n" + "In the case that the semantics for the operator are not trivial, it would be appreciated " + "to also include a proposal for the semantics." + ) + warnings.warn(msg) + return NotImplemented + if not is_masked_tensor(self): + raise TypeError("Input to reduce_dim must be a MaskedTensor") + + masked_fn = _get_masked_fn(fn) + data = self.get_data() + mask = self.get_mask() + if fn == "all": + result_data = masked_fn(data, dim=dim, keepdim=keepdim, mask=mask) + else: + result_data = masked_fn( + self, dim=dim, keepdim=keepdim, dtype=dtype, mask=self.get_mask() + ) + return as_masked_tensor(result_data, _multidim_any(mask, dim, keepdim)) + + return reduce_dim + + +def _torch_reduce(fn): + def reduce_fn(*args, **kwargs): + if len(args) == 1 and len(kwargs) == 0: + return _torch_reduce_all(fn)(args[0]) + return _torch_reduce_dim(fn)(*args, **kwargs) + + return reduce_fn + + +def _reduce_dim_args(input, dim, keepdim=False, dtype=None): + return input, dim, keepdim, dtype + + +def _torch_grad_reduce(fn): + def grad_reduce(*args, **kwargs): + if len(args) == 1 and len(kwargs) == 0: + return _torch_reduce_all(fn)(args[0]) + # TODO: autograd.Function doesn't support kwarg + input, dim, keepdim, dtype = _reduce_dim_args(*args, **kwargs) + return _torch_reduce_dim(fn)(input, dim, keepdim, dtype) + + return grad_reduce + + +REDUCE_NAMES = [ + "sum", + "mean", + "amin", + "amax", + "argmin", + "argmax", + "prod", + "all", + "norm", + "var", + "std", +] + +NATIVE_REDUCE_MAP = { + getattr(torch.ops.aten, name): _torch_reduce(name) for name in REDUCE_NAMES +} +TORCH_REDUCE_MAP = { + getattr(torch, name): _torch_grad_reduce(name) for name in REDUCE_NAMES +} +TENSOR_REDUCE_MAP = { + getattr(torch.Tensor, name): _torch_grad_reduce(name) for name in REDUCE_NAMES +} + +NATIVE_REDUCE_FNS = list(NATIVE_REDUCE_MAP.keys()) +TORCH_REDUCE_FNS = list(TORCH_REDUCE_MAP.keys()) +TENSOR_REDUCE_FNS = list(TENSOR_REDUCE_MAP.keys()) + + +def _is_reduction(fn): + return fn in NATIVE_REDUCE_MAP or fn in TORCH_REDUCE_MAP or fn in TENSOR_REDUCE_MAP + + +def _apply_reduction(fn, *args, **kwargs): + if fn in NATIVE_REDUCE_MAP: + return NATIVE_REDUCE_MAP[fn](*args, **kwargs) + if fn in TORCH_REDUCE_MAP: + return TORCH_REDUCE_MAP[fn](*args, **kwargs) + if fn in TENSOR_REDUCE_MAP: + return TENSOR_REDUCE_MAP[fn](*args, **kwargs) + return NotImplemented diff --git a/lib/python3.12/site-packages/torch/masked/maskedtensor/unary.py b/lib/python3.12/site-packages/torch/masked/maskedtensor/unary.py new file mode 100644 index 0000000000000000000000000000000000000000..e04ee6e810a7418829b68323097612391017b14e --- /dev/null +++ b/lib/python3.12/site-packages/torch/masked/maskedtensor/unary.py @@ -0,0 +1,194 @@ +# mypy: allow-untyped-defs +# Copyright (c) Meta Platforms, Inc. and affiliates + +import torch + +from .core import _map_mt_args_kwargs, _wrap_result + + +__all__ = [] # type: ignore[var-annotated] + + +UNARY_NAMES = [ + "abs", + "absolute", + "acos", + "arccos", + "acosh", + "arccosh", + "angle", + "asin", + "arcsin", + "asinh", + "arcsinh", + "atan", + "arctan", + "atanh", + "arctanh", + "bitwise_not", + "ceil", + "clamp", + "clip", + "conj_physical", + "cos", + "cosh", + "deg2rad", + "digamma", + "erf", + "erfc", + "erfinv", + "exp", + "exp2", + "expm1", + "fix", + "floor", + "frac", + "lgamma", + "log", + "log10", + "log1p", + "log2", + "logit", + "i0", + "isnan", + "nan_to_num", + "neg", + "negative", + "positive", + "pow", + "rad2deg", + "reciprocal", + "round", + "rsqrt", + "sigmoid", + "sign", + "sgn", + "signbit", + "sin", + "sinc", + "sinh", + "sqrt", + "square", + "tan", + "tanh", + "trunc", +] + +INPLACE_UNARY_NAMES = [ + n + "_" + for n in (list(set(UNARY_NAMES) - {"angle", "positive", "signbit", "isnan"})) +] + +# Explicitly tracking functions we know are currently not supported +# This might be due to missing code gen or because of complex semantics +UNARY_NAMES_UNSUPPORTED = [ + "atan2", + "arctan2", + "bitwise_left_shift", + "bitwise_right_shift", + "copysign", + "float_power", + "fmod", + "frexp", + "gradient", + "imag", + "ldexp", + "lerp", + "logical_not", + "hypot", + "igamma", + "igammac", + "mvlgamma", + "nextafter", + "polygamma", + "real", + "remainder", + "true_divide", + "xlogy", +] + + +def _unary_helper(fn, args, kwargs, inplace): + if len(kwargs) != 0: + raise ValueError( + "MaskedTensor unary ops require that len(kwargs) == 0. " + "If you need support for this, please open an issue on Github." + ) + for a in args[1:]: + if torch.is_tensor(a): + raise TypeError( + "MaskedTensor unary ops do not support additional Tensor arguments" + ) + + mask_args, _mask_kwargs = _map_mt_args_kwargs( + args, kwargs, lambda x: x._masked_mask + ) + data_args, _data_kwargs = _map_mt_args_kwargs( + args, kwargs, lambda x: x._masked_data + ) + + if args[0].layout == torch.sparse_coo: + data_args[0] = data_args[0].coalesce() + s = data_args[0].size() + i = data_args[0].indices() + data_args[0] = data_args[0].coalesce().values() + v = fn(*data_args) + result_data = torch.sparse_coo_tensor(i, v, size=s) + + elif args[0].layout == torch.sparse_csr: + crow = data_args[0].crow_indices() + col = data_args[0].col_indices() + data_args[0] = data_args[0].values() + v = fn(*data_args) + result_data = torch.sparse_csr_tensor(crow, col, v) + + else: + result_data = fn(*data_args) + + if inplace: + args[0]._set_data_mask(result_data, mask_args[0]) + return args[0] + else: + return _wrap_result(result_data, mask_args[0]) + + +def _torch_unary(fn_name): + fn = getattr(torch.ops.aten, fn_name) + + def unary_fn(*args, **kwargs): + return _unary_helper(fn, args, kwargs, inplace=False) + + return unary_fn + + +def _torch_inplace_unary(fn_name): + fn = getattr(torch.ops.aten, fn_name) + + def unary_fn(*args, **kwargs): + return _unary_helper(fn, args, kwargs, inplace=True) + + return unary_fn + + +NATIVE_UNARY_MAP = { + getattr(torch.ops.aten, name): _torch_unary(name) for name in UNARY_NAMES +} +NATIVE_INPLACE_UNARY_MAP = { + getattr(torch.ops.aten, name): _torch_inplace_unary(name) + for name in INPLACE_UNARY_NAMES +} + +NATIVE_UNARY_FNS = list(NATIVE_UNARY_MAP.keys()) +NATIVE_INPLACE_UNARY_FNS = list(NATIVE_INPLACE_UNARY_MAP.keys()) + + +def _is_native_unary(fn): + return fn in NATIVE_UNARY_FNS or fn in NATIVE_INPLACE_UNARY_FNS + + +def _apply_native_unary(fn, *args, **kwargs): + if fn in NATIVE_UNARY_FNS: + return NATIVE_UNARY_MAP[fn](*args, **kwargs) + if fn in NATIVE_INPLACE_UNARY_FNS: + return NATIVE_INPLACE_UNARY_MAP[fn](*args, **kwargs) + return NotImplemented diff --git a/lib/python3.12/site-packages/torch/mps/__init__.py b/lib/python3.12/site-packages/torch/mps/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..94c56fc4d7da502f05a49df1c13d0759d8ed5573 --- /dev/null +++ b/lib/python3.12/site-packages/torch/mps/__init__.py @@ -0,0 +1,183 @@ +# mypy: allow-untyped-defs +r""" +This package enables an interface for accessing MPS (Metal Performance Shaders) backend in Python. +Metal is Apple's API for programming metal GPU (graphics processor unit). Using MPS means that increased +performance can be achieved, by running work on the metal GPU(s). +See https://developer.apple.com/documentation/metalperformanceshaders for more details. +""" +from typing import Union + +import torch +from torch import Tensor + + +_is_in_bad_fork = getattr(torch._C, "_mps_is_in_bad_fork", lambda: False) +_default_mps_generator: torch._C.Generator = None # type: ignore[assignment] + + +# local helper function (not public or exported) +def _get_default_mps_generator() -> torch._C.Generator: + global _default_mps_generator + if _default_mps_generator is None: + _default_mps_generator = torch._C._mps_get_default_generator() + return _default_mps_generator + + +def device_count() -> int: + r"""Returns the number of available MPS devices.""" + return int(torch._C._has_mps and torch._C._mps_is_available()) + + +def synchronize() -> None: + r"""Waits for all kernels in all streams on a MPS device to complete.""" + return torch._C._mps_deviceSynchronize() + + +def get_rng_state(device: Union[int, str, torch.device] = "mps") -> Tensor: + r"""Returns the random number generator state as a ByteTensor. + + Args: + device (torch.device or int, optional): The device to return the RNG state of. + Default: ``'mps'`` (i.e., ``torch.device('mps')``, the current MPS device). + """ + return _get_default_mps_generator().get_state() + + +def set_rng_state( + new_state: Tensor, device: Union[int, str, torch.device] = "mps" +) -> None: + r"""Sets the random number generator state. + + Args: + new_state (torch.ByteTensor): The desired state + device (torch.device or int, optional): The device to set the RNG state. + Default: ``'mps'`` (i.e., ``torch.device('mps')``, the current MPS device). + """ + new_state_copy = new_state.clone(memory_format=torch.contiguous_format) + _get_default_mps_generator().set_state(new_state_copy) + + +def manual_seed(seed: int) -> None: + r"""Sets the seed for generating random numbers. + + Args: + seed (int): The desired seed. + """ + # the torch.mps.manual_seed() can be called from the global + # torch.manual_seed() in torch/random.py. So we need to make + # sure mps is available (otherwise we just return without + # erroring out) + if not torch._C._has_mps: + return + seed = int(seed) + _get_default_mps_generator().manual_seed(seed) + + +def seed() -> None: + r"""Sets the seed for generating random numbers to a random number.""" + _get_default_mps_generator().seed() + + +def empty_cache() -> None: + r"""Releases all unoccupied cached memory currently held by the caching + allocator so that those can be used in other GPU applications. + """ + torch._C._mps_emptyCache() + + +def set_per_process_memory_fraction(fraction) -> None: + r"""Set memory fraction for limiting process's memory allocation on MPS device. + The allowed value equals the fraction multiplied by recommended maximum device memory + (obtained from Metal API device.recommendedMaxWorkingSetSize). + If trying to allocate more than the allowed value in a process, it will raise an out of + memory error in allocator. + + Args: + fraction(float): Range: 0~2. Allowed memory equals total_memory * fraction. + + .. note:: + Passing 0 to fraction means unlimited allocations + (may cause system failure if out of memory). + Passing fraction greater than 1.0 allows limits beyond the value + returned from device.recommendedMaxWorkingSetSize. + """ + + if not isinstance(fraction, float): + raise TypeError("Invalid type for fraction argument, must be `float`") + if fraction < 0 or fraction > 2: + raise ValueError(f"Invalid fraction value: {fraction}. Allowed range: 0~2") + + torch._C._mps_setMemoryFraction(fraction) + + +def current_allocated_memory() -> int: + r"""Returns the current GPU memory occupied by tensors in bytes. + + .. note:: + The returned size does not include cached allocations in + memory pools of MPSAllocator. + """ + return torch._C._mps_currentAllocatedMemory() + + +def driver_allocated_memory() -> int: + r"""Returns total GPU memory allocated by Metal driver for the process in bytes. + + .. note:: + The returned size includes cached allocations in MPSAllocator pools + as well as allocations from MPS/MPSGraph frameworks. + """ + return torch._C._mps_driverAllocatedMemory() + + +def recommended_max_memory() -> int: + r"""Returns recommended max Working set size for GPU memory in bytes. + + .. note:: + Recommended max working set size for Metal. + returned from device.recommendedMaxWorkingSetSize. + """ + return torch._C._mps_recommendedMaxMemory() + + +def _compile_shader(source: str): + r"""Compiles compute shader from source and allows one to invoke kernels + defined there from the comfort of Python runtime + Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_MPS) + >>> lib = torch.mps._compile_shader( + ... "kernel void full(device float* out, constant float& val, uint idx [[thread_position_in_grid]]) { out[idx] = val; }" + ... ) + >>> x = torch.zeros(16, device="mps") + >>> lib.full(x, 3.14) + """ + if not hasattr(torch._C, "_mps_compileShader"): + raise RuntimeError("MPS is not available") + return torch._C._mps_compileShader(source) + + +def is_available() -> bool: + return device_count() > 0 + + +from . import profiler +from .event import Event + + +__all__ = [ + "device_count", + "get_rng_state", + "manual_seed", + "seed", + "set_rng_state", + "synchronize", + "empty_cache", + "set_per_process_memory_fraction", + "current_allocated_memory", + "driver_allocated_memory", + "Event", + "profiler", + "recommended_max_memory", + "is_available", +] diff --git a/lib/python3.12/site-packages/torch/mps/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/mps/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3af9e13b035b062bf804f10720f2ddc4cbbb07d1 Binary files /dev/null and b/lib/python3.12/site-packages/torch/mps/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/mps/__pycache__/event.cpython-312.pyc b/lib/python3.12/site-packages/torch/mps/__pycache__/event.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25a7b037c85c67cc94dfa266e52b29444e569706 Binary files /dev/null and b/lib/python3.12/site-packages/torch/mps/__pycache__/event.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/mps/__pycache__/profiler.cpython-312.pyc b/lib/python3.12/site-packages/torch/mps/__pycache__/profiler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..709922a12540c9645a19ca7f14f519a40af471f7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/mps/__pycache__/profiler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/mps/event.py b/lib/python3.12/site-packages/torch/mps/event.py new file mode 100644 index 0000000000000000000000000000000000000000..d619c027480c3ad6c52744afa76f35ff4cba64c0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/mps/event.py @@ -0,0 +1,46 @@ +# mypy: allow-untyped-defs +import torch + + +class Event: + r"""Wrapper around an MPS event. + + MPS events are synchronization markers that can be used to monitor the + device's progress, to accurately measure timing, and to synchronize MPS streams. + + Args: + enable_timing (bool, optional): indicates if the event should measure time + (default: ``False``) + """ + + def __init__(self, enable_timing=False): + self.__eventId = torch._C._mps_acquireEvent(enable_timing) + + def __del__(self): + # checks if torch._C is already destroyed + if hasattr(torch._C, "_mps_releaseEvent") and self.__eventId > 0: + torch._C._mps_releaseEvent(self.__eventId) + + def record(self): + r"""Records the event in the default stream.""" + torch._C._mps_recordEvent(self.__eventId) + + def wait(self): + r"""Makes all future work submitted to the default stream wait for this event.""" + torch._C._mps_waitForEvent(self.__eventId) + + def query(self): + r"""Returns True if all work currently captured by event has completed.""" + return torch._C._mps_queryEvent(self.__eventId) + + def synchronize(self): + r"""Waits until the completion of all work currently captured in this event. + This prevents the CPU thread from proceeding until the event completes. + """ + torch._C._mps_synchronizeEvent(self.__eventId) + + def elapsed_time(self, end_event): + r"""Returns the time elapsed in milliseconds after the event was + recorded and before the end_event was recorded. + """ + return torch._C._mps_elapsedTimeOfEvents(self.__eventId, end_event.__eventId) diff --git a/lib/python3.12/site-packages/torch/mps/profiler.py b/lib/python3.12/site-packages/torch/mps/profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..4dcd86b30ecbf765cb13c674e602bdf21058f668 --- /dev/null +++ b/lib/python3.12/site-packages/torch/mps/profiler.py @@ -0,0 +1,61 @@ +# mypy: allow-untyped-defs +import contextlib + +import torch + + +__all__ = ["start", "stop", "profile"] + + +def start(mode: str = "interval", wait_until_completed: bool = False) -> None: + r"""Start OS Signpost tracing from MPS backend. + + The generated OS Signposts could be recorded and viewed in + XCode Instruments Logging tool. + + Args: + mode(str): OS Signpost tracing mode could be "interval", "event", + or both "interval,event". + The interval mode traces the duration of execution of the operations, + whereas event mode marks the completion of executions. + See document `Recording Performance Data`_ for more info. + wait_until_completed(bool): Waits until the MPS Stream complete + executing each encoded GPU operation. This helps generating single + dispatches on the trace's timeline. + Note that enabling this option would affect the performance negatively. + + .. _Recording Performance Data: + https://developer.apple.com/documentation/os/logging/recording_performance_data + """ + mode_normalized = mode.lower().replace(" ", "") + torch._C._mps_profilerStartTrace(mode_normalized, wait_until_completed) + + +def stop(): + r"""Stops generating OS Signpost tracing from MPS backend.""" + torch._C._mps_profilerStopTrace() + + +@contextlib.contextmanager +def profile(mode: str = "interval", wait_until_completed: bool = False): + r"""Context Manager to enabling generating OS Signpost tracing from MPS backend. + + Args: + mode(str): OS Signpost tracing mode could be "interval", "event", + or both "interval,event". + The interval mode traces the duration of execution of the operations, + whereas event mode marks the completion of executions. + See document `Recording Performance Data`_ for more info. + wait_until_completed(bool): Waits until the MPS Stream complete + executing each encoded GPU operation. This helps generating single + dispatches on the trace's timeline. + Note that enabling this option would affect the performance negatively. + + .. _Recording Performance Data: + https://developer.apple.com/documentation/os/logging/recording_performance_data + """ + try: + start(mode, wait_until_completed) + yield + finally: + stop() diff --git a/lib/python3.12/site-packages/torch/mtia/__init__.py b/lib/python3.12/site-packages/torch/mtia/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..170d9f7cb6777dc8b4ca611bbb08c646f4b4e37b --- /dev/null +++ b/lib/python3.12/site-packages/torch/mtia/__init__.py @@ -0,0 +1,339 @@ +# mypy: allow-untyped-defs +r""" +This package enables an interface for accessing MTIA backend in python +""" + +import threading +import warnings +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import torch +from torch import device as _device, Tensor +from torch._utils import _dummy_type, _LazySeedTracker, classproperty +from torch.types import Device + +from ._utils import _get_device_index + + +_device_t = Union[_device, str, int] + +# torch.mtia.Event/Stream is alias of torch.Event/Stream +Event = torch.Event +Stream = torch.Stream + +_initialized = False +_queued_calls: List[ + Tuple[Callable[[], None], List[str]] +] = [] # don't invoke these until initialization occurs +_tls = threading.local() +_initialization_lock = threading.Lock() +_lazy_seed_tracker = _LazySeedTracker() + + +def init(): + _lazy_init() + + +def is_initialized(): + r"""Return whether PyTorch's MTIA state has been initialized.""" + return _initialized and not _is_in_bad_fork() + + +def _is_in_bad_fork() -> bool: + return torch._C._mtia_isInBadFork() + + +def _lazy_init() -> None: + global _initialized, _queued_calls + if is_initialized() or hasattr(_tls, "is_initializing"): + return + with _initialization_lock: + # We be double-checking locking, boys! This is OK because + # the above test was GIL protected anyway. The inner test + # is for when a thread blocked on some other thread which was + # doing the initialization; when they get the lock, they will + # find there is nothing left to do. + if is_initialized(): + return + # It is important to prevent other threads from entering _lazy_init + # immediately, while we are still guaranteed to have the GIL, because some + # of the C calls we make below will release the GIL + if _is_in_bad_fork(): + raise RuntimeError( + "Cannot re-initialize MTIA in forked subprocess. To use MTIA with " + "multiprocessing, you must use the 'spawn' start method" + ) + if not _is_compiled(): + raise AssertionError( + "Torch not compiled with MTIA enabled. " + "Ensure you have `import mtia.host_runtime.torch_mtia` in your python " + "src file and include `//mtia/host_runtime/torch_mtia:torch_mtia` as " + "your target dependency!" + ) + + torch._C._mtia_init() + # Some of the queued calls may reentrantly call _lazy_init(); + # we need to just return without initializing in that case. + # However, we must not let any *other* threads in! + _tls.is_initializing = True + + _queued_calls.extend(calls for calls in _lazy_seed_tracker.get_calls() if calls) + + try: + for queued_call, orig_traceback in _queued_calls: + try: + queued_call() + except Exception as e: + msg = ( + f"MTIA call failed lazily at initialization with error: {str(e)}\n\n" + f"MTIA call was originally invoked at:\n\n{''.join(orig_traceback)}" + ) + raise DeferredMtiaCallError(msg) from e + finally: + delattr(_tls, "is_initializing") + _initialized = True + + +class DeferredMtiaCallError(Exception): + pass + + +def _is_compiled() -> bool: + r"""Return true if compiled with MTIA support.""" + return torch._C._mtia_isBuilt() + + +def is_available() -> bool: + r"""Return true if MTIA device is available""" + if not _is_compiled(): + return False + # MTIA has to init devices first to know if there is any devices available. + return device_count() > 0 + + +def synchronize(device: Optional[_device_t] = None) -> None: + r"""Waits for all jobs in all streams on a MTIA device to complete.""" + with torch.mtia.device(device): + return torch._C._mtia_deviceSynchronize() + + +def device_count() -> int: + r"""Return the number of MTIA devices available.""" + return torch._C._accelerator_hooks_device_count() + + +def current_device() -> int: + r"""Return the index of a currently selected device.""" + return torch._C._accelerator_hooks_get_current_device() + + +def current_stream(device: Optional[_device_t] = None) -> Stream: + r"""Return the currently selected :class:`Stream` for a given device. + + Args: + device (torch.device or int, optional): selected device. Returns + the currently selected :class:`Stream` for the current device, given + by :func:`~torch.mtia.current_device`, if :attr:`device` is ``None`` + (default). + """ + return torch._C._mtia_getCurrentStream(_get_device_index(device, optional=True)) + + +def default_stream(device: Optional[_device_t] = None) -> Stream: + r"""Return the default :class:`Stream` for a given device. + + Args: + device (torch.device or int, optional): selected device. Returns + the default :class:`Stream` for the current device, given by + :func:`~torch.mtia.current_device`, if :attr:`device` is ``None`` + (default). + """ + return torch._C._mtia_getDefaultStream(_get_device_index(device, optional=True)) + + +def get_device_capability(device: Optional[_device_t] = None) -> Tuple[int, int]: + r"""Return capability of a given device as a tuple of (major version, minor version). + + Args: + device (torch.device or int, optional) selected device. Returns + statistics for the current device, given by current_device(), + if device is None (default). + """ + return torch._C._mtia_getDeviceCapability(_get_device_index(device, optional=True)) + + +def empty_cache() -> None: + r"""Empty the MTIA device cache.""" + return torch._C._mtia_emptyCache() + + +def set_stream(stream: Stream): + r"""Set the current stream.This is a wrapper API to set the stream. + Usage of this function is discouraged in favor of the ``stream`` + context manager. + + Args: + stream (Stream): selected stream. This function is a no-op + if this argument is ``None``. + """ + if stream is None: + return + torch._C._mtia_setCurrentStream(stream) + + +def set_device(device: _device_t) -> None: + r"""Set the current device. + + Args: + device (torch.device or int): selected device. This function is a no-op + if this argument is negative. + """ + device = _get_device_index(device) + if device >= 0: + torch._C._accelerator_hooks_set_current_device(device) + + +class device: + r"""Context-manager that changes the selected device. + + Args: + device (torch.device or int): device index to select. It's a no-op if + this argument is a negative integer or ``None``. + """ + + def __init__(self, device: Any): + self.idx = _get_device_index(device, optional=True) + self.prev_idx = -1 + + def __enter__(self): + self.prev_idx = torch._C._accelerator_hooks_maybe_exchange_device(self.idx) + + def __exit__(self, type: Any, value: Any, traceback: Any): + self.idx = torch._C._accelerator_hooks_maybe_exchange_device(self.prev_idx) + return False + + +class StreamContext: + r"""Context-manager that selects a given stream. + + All MTIA kernels queued within its context will be enqueued on a selected + stream. + + Args: + Stream (Stream): selected stream. This manager is a no-op if it's + ``None``. + .. note:: Streams are per-device. + """ + + cur_stream: Optional["torch.mtia.Stream"] + + def __init__(self, stream: Optional["torch.mtia.Stream"]): + self.cur_stream = None + self.stream = stream + self.idx = _get_device_index(None, True) + if not torch.jit.is_scripting(): + if self.idx is None: + self.idx = -1 + + self.src_prev_stream = ( + None if not torch.jit.is_scripting() else torch.mtia.default_stream(None) + ) + self.dst_prev_stream = ( + None if not torch.jit.is_scripting() else torch.mtia.default_stream(None) + ) + + def __enter__(self): + # Local cur_stream variable for type refinement + cur_stream = self.stream + # Return if stream is None or MTIA device not available + if cur_stream is None or self.idx == -1: + return + self.src_prev_stream = torch.mtia.current_stream(None) + + # If the stream is not on the current device, then + # set the current stream on the device + if self.src_prev_stream.device != cur_stream.device: + with device(cur_stream.device): + self.dst_prev_stream = torch.mtia.current_stream(cur_stream.device) + torch.mtia.set_stream(cur_stream) + + def __exit__(self, type: Any, value: Any, traceback: Any): + # Local cur_stream variable for type refinement + cur_stream = self.stream + # If stream is None or no MTIA device available, return + if cur_stream is None or self.idx == -1: + return + + # Reset the stream on the original device + # and destination device + if self.src_prev_stream.device != cur_stream.device: # type: ignore[union-attr] + torch.mtia.set_stream(self.dst_prev_stream) # type: ignore[arg-type] + torch.mtia.set_stream(self.src_prev_stream) # type: ignore[arg-type] + + +def stream(stream: Optional["torch.mtia.Stream"]) -> StreamContext: + r"""Wrap around the Context-manager StreamContext that selects a given stream. + + Arguments: + stream (Stream): selected stream. This manager is a no-op if it's + ``None``. + ..Note:: In eager mode stream is of type Stream class while in JIT it doesn't support torch.mtia.stream + """ + return StreamContext(stream) + + +def get_rng_state(device: Union[int, str, torch.device] = "mtia") -> Tensor: + r"""Returns the random number generator state as a ByteTensor. + + Args: + device (torch.device or int, optional): The device to return the RNG state of. + Default: ``'mtia'`` (i.e., ``torch.device('mtia')``, the current mtia device). + """ + warnings.warn( + "get_rng_state is not implemented in torch.mtia", + UserWarning, + stacklevel=2, + ) + return torch.zeros([1], dtype=torch.uint8, device=device) + + +def set_rng_state( + new_state: Tensor, device: Union[int, str, torch.device] = "mtia" +) -> None: + r"""Sets the random number generator state. + + Args: + new_state (torch.ByteTensor): The desired state + device (torch.device or int, optional): The device to set the RNG state. + Default: ``'mtia'`` (i.e., ``torch.device('mtia')``, the current mtia device). + """ + warnings.warn( + "set_rng_state is not implemented in torch.mtia", + UserWarning, + stacklevel=2, + ) + + +from .memory import * # noqa: F403 + + +__all__ = [ + "init", + "is_available", + "is_initialized", + "synchronize", + "device_count", + "current_device", + "current_stream", + "default_stream", + "memory_stats", + "max_memory_allocated", + "get_device_capability", + "empty_cache", + "set_device", + "set_stream", + "stream", + "device", + "set_rng_state", + "get_rng_state", +] diff --git a/lib/python3.12/site-packages/torch/mtia/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/mtia/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..afccbfc3a9d587a7aaaf37d82b1174f95ceb8b46 Binary files /dev/null and b/lib/python3.12/site-packages/torch/mtia/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/mtia/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/mtia/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff196e617a9bac0f2e5cfa84b92c63bd8af0910f Binary files /dev/null and b/lib/python3.12/site-packages/torch/mtia/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/mtia/__pycache__/memory.cpython-312.pyc b/lib/python3.12/site-packages/torch/mtia/__pycache__/memory.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c4439e9531c33b49a62f3806e900862c0b1c8dbb Binary files /dev/null and b/lib/python3.12/site-packages/torch/mtia/__pycache__/memory.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/mtia/_utils.py b/lib/python3.12/site-packages/torch/mtia/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..090e26f321232f9687c2b348ac602dbb6699b03f --- /dev/null +++ b/lib/python3.12/site-packages/torch/mtia/_utils.py @@ -0,0 +1,38 @@ +from typing import Any + +import torch + +# The _get_device_index has been moved to torch.utils._get_device_index +from torch._utils import _get_device_index as _torch_get_device_index + + +def _get_device_index( + device: Any, optional: bool = False, allow_cpu: bool = False +) -> int: + r"""Get the device index from :attr:`device`, which can be a torch.device object, a Python integer, or ``None``. + + If :attr:`device` is a torch.device object, returns the device index if it + is a MTIA device. Note that for a MTIA device without a specified index, + i.e., ``torch.device('mtia')``, this will return the current default MTIA + device if :attr:`optional` is ``True``. If :attr:`allow_cpu` is ``True``, + CPU devices will be accepted and ``-1`` will be returned in this case. + + If :attr:`device` is a Python integer, it is returned as is. + + If :attr:`device` is ``None``, this will return the current default MTIA + device if :attr:`optional` is ``True``. + """ + if isinstance(device, int): + return device + if isinstance(device, str): + device = torch.device(device) + if isinstance(device, torch.device): + if allow_cpu: + if device.type not in ["mtia", "cpu"]: + raise ValueError(f"Expected a mtia or cpu device, but got: {device}") + elif device.type != "mtia": + raise ValueError(f"Expected a mtia device, but got: {device}") + if not torch.jit.is_scripting(): + if isinstance(device, torch.mtia.device): + return device.idx + return _torch_get_device_index(device, optional, allow_cpu) diff --git a/lib/python3.12/site-packages/torch/mtia/memory.py b/lib/python3.12/site-packages/torch/mtia/memory.py new file mode 100644 index 0000000000000000000000000000000000000000..0686c71c7e43fe6ba8831bc4864596d79ef84cb5 --- /dev/null +++ b/lib/python3.12/site-packages/torch/mtia/memory.py @@ -0,0 +1,41 @@ +# pyre-strict + +r"""This package adds support for device memory management implemented in MTIA.""" + +from typing import Any, Dict, Optional + +import torch + +from . import _device_t, is_initialized +from ._utils import _get_device_index + + +def memory_stats(device: Optional[_device_t] = None) -> Dict[str, Any]: + r"""Return a dictionary of MTIA memory allocator statistics for a given device. + + Args: + device (torch.device, str, or int, optional) selected device. Returns + statistics for the current device, given by current_device(), + if device is None (default). + """ + if not is_initialized(): + return {} + return torch._C._mtia_memoryStats(_get_device_index(device, optional=True)) + + +def max_memory_allocated(device: Optional[_device_t] = None) -> int: + r"""Return the maximum memory allocated in bytes for a given device. + + Args: + device (torch.device or int, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.mtia.current_device`, + if :attr:`device` is ``None`` (default). + """ + + return memory_stats(device=device).get("allocated_bytes.all.peak", 0) + + +__all__ = [ + "memory_stats", + "max_memory_allocated", +] diff --git a/lib/python3.12/site-packages/torch/testing/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b4cbda089100a502e8bf8b32bb99cc4e64751ad Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/__pycache__/_comparison.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/__pycache__/_comparison.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d95aaeb6fa4a61a0cf142384fd6b37b23d9c08ff Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/__pycache__/_comparison.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/__pycache__/_creation.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/__pycache__/_creation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff395b5482b13c11f451bff4d805430b38f802d8 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/__pycache__/_creation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7114327cc4d2e6139bc0735d5bc96a84a62fc878 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7bc09dd398666bf9969d2f2bead01e139d0518d3 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/checkpoint_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/checkpoint_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33c99eb50dd9ca8a8caa525930d90ac5d7f79d4e Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/checkpoint_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/common_state_dict.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/common_state_dict.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b9c86b1b4e382627432ddf58f69a0e3a8dd50eeb Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/common_state_dict.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/ddp_under_dist_autograd_test.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/ddp_under_dist_autograd_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4362215e1f6389b3376923643a000e952181c0ad Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/ddp_under_dist_autograd_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/distributed_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/distributed_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5558b8f527ac4eee38c352258aa6caf73b5cd02a Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/distributed_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/fake_pg.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/fake_pg.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21f4d9543c12cdeda421bb35f8cd97b4f6cbf3d6 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/fake_pg.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/multi_threaded_pg.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/multi_threaded_pg.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e304706771581156aa239a99d57a4312e0f76aa8 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/multi_threaded_pg.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/rpc_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/rpc_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92deba683330c9a4a78375aeed3842579b68ca3b Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/__pycache__/rpc_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd083b4e2ad974c429eee50b1231ff79bc1df090 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/__pycache__/test_common.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/__pycache__/test_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd5156f5f6be5e8be96261fe6ed72768be9b3c63 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/__pycache__/test_common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__init__.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..303d06063fda26d56bf999b3041f779492f16af3 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__init__.py @@ -0,0 +1,98 @@ +# mypy: allow-untyped-defs + +import sys +from functools import wraps, partial + +import torch +import torch.distributed as dist +from torch.distributed import rpc +from torch.testing._internal.common_distributed import ( + MultiProcessTestCase, + TEST_SKIPS, + tp_transports, +) + +TEST_GPU_NUM = 4 + +class ShardedTensorTestBase(MultiProcessTestCase): + @property + def world_size(self): + return TEST_GPU_NUM + + def init_pg(self, backend="nccl"): + if backend not in ["nccl", "gloo", "mpi"]: + raise RuntimeError(f"Backend {backend} not supported!") + + dist.init_process_group( + backend=backend, + world_size=self.world_size, + rank=self.rank, + init_method=f"file://{self.file_name}", + ) + + # set device for nccl pg for collectives + if backend == "nccl": + torch.cuda.set_device(self.rank) + + + def init_rpc(self): + rpc_backend_options = rpc.TensorPipeRpcBackendOptions(_transports=tp_transports()) + rpc_backend_options.init_method = f"file://{self.file_name}" + for rank in range(self.world_size): + rpc_backend_options.set_device_map( + f"worker{rank}", {rank: self.rank, self.rank: rank} + ) + + rpc.init_rpc( + name="worker%d" % self.rank, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + + def init_comms(self, init_rpc=True, backend="nccl"): + if init_rpc: + self.init_rpc() + self.init_pg(backend=backend) + + def destroy_comms(self, destroy_rpc=True): + # Wait for all ranks to reach here before starting shutdown. + dist.barrier() + + if destroy_rpc: + rpc.shutdown() + dist.destroy_process_group() + + def setUp(self) -> None: + super().setUp() + self._spawn_processes() + + def assert_sharded_tensor_equal(self, st1, st2): + st1_local_shards = st1.local_shards() + st2_local_shards = st2.local_shards() + self.assertEqual(len(st1_local_shards), len(st2_local_shards)) + for i, st1_local_shard in enumerate(st1_local_shards): + self.assertEqual(st1_local_shard.tensor, st2_local_shards[i].tensor) + self.assertEqual(st1_local_shard.metadata, st2_local_shards[i].metadata) + + self.assertEqual(st1.metadata(), st2.metadata()) + self.assertEqual(st1.sharding_spec(), st2.sharding_spec()) + self.assertEqual(len(st1.remote_shards()), len(st2.remote_shards())) + +# wrapper to initialize comms (processgroup + rpc) +def with_comms(func=None, init_rpc=True, backend="nccl"): + if func is None: + return partial( + with_comms, + init_rpc=init_rpc, + backend=backend, + ) + + @wraps(func) + def wrapper(self, *args, **kwargs): + if backend == "nccl" and torch.cuda.device_count() < self.world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code) + self.init_comms(init_rpc=init_rpc, backend=backend) + func(self, *args, **kwargs) + self.destroy_comms(destroy_rpc=init_rpc) + return wrapper diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9247d24dfd0d3f64b0aa2ac263daef415fe07300 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/_test_ops_common.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/_test_ops_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e7759546e34dd9f454b082ca7c994392ba26d94 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/_test_ops_common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/_test_st_common.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/_test_st_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91e9b6442197f46904e61ee4c4da121f00d4be4b Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/_test_st_common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_ops_common.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_ops_common.py new file mode 100644 index 0000000000000000000000000000000000000000..398b2fd8a36aa2b43e67ae0d161ab0df1c1d51d4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_ops_common.py @@ -0,0 +1,136 @@ +# mypy: allow-untyped-defs + +import builtins + +import torch +from torch.distributed._shard.sharding_spec import ( + ChunkShardingSpec, + EnumerableShardingSpec, + ShardMetadata, +) +from torch.distributed._shard.sharding_spec._internals import ( + get_chunked_dim_size, + get_split_size, +) + + +def generate_chunk_sharding_specs_for_test(sharding_dim): + return [ + ChunkShardingSpec( + dim=sharding_dim, + placements=[ + "rank:0/cuda:0", + "rank:1/cuda:1", + "rank:2/cuda:2", + "rank:3/cuda:3", + ], + ), + # Test different ordering. (Case 1) + ChunkShardingSpec( + dim=sharding_dim, + placements=[ + "rank:2/cuda:2", + "rank:3/cuda:3", + "rank:0/cuda:0", + "rank:1/cuda:1", + ], + ), + # Test different ordering. (Case 2) + ChunkShardingSpec( + dim=sharding_dim, + placements=[ + "rank:3/cuda:3", + "rank:0/cuda:0", + "rank:1/cuda:1", + "rank:2/cuda:2", + ], + ), + ] + + +def generate_enumerable_sharding_specs_for_test(): + return [ + EnumerableShardingSpec( + [ + ShardMetadata( + shard_offsets=[0, 0], + shard_sizes=[5, 5], + placement="rank:0/cuda:0", + ), + ShardMetadata( + shard_offsets=[5, 0], + shard_sizes=[5, 5], + placement="rank:1/cuda:1", + ), + ShardMetadata( + shard_offsets=[0, 5], + shard_sizes=[5, 5], + placement="rank:2/cuda:2", + ), + ShardMetadata( + shard_offsets=[5, 5], + shard_sizes=[5, 5], + placement="rank:3/cuda:3", + ), + ] + ) + ] + + +def generate_local_weight_sharding_params_for_test( + local_weight, sharded_dim, gpu_num, spec, rank +): + """ + Shard the local weight based the given spec, so we can compare against + the one from sharded tensor. + + Args: + local_weight: weight matrix to be sharded. + sharded_dim: The dimension which we shard on. + gpu_num: number of ranks. + spec: sharding spec. + rank: # of cuda process. + + Returns: + start_pos: start position of sharded weight on the given rank. + chunk_size: chunk size of sharded weight on the given rank. + """ + sharding_dim_size = local_weight.size(sharded_dim) + split_size = get_split_size(sharding_dim_size, gpu_num) + current_offsets = 0 + start_pos = current_offsets + for idx, placement in enumerate(spec.placements): + chunk_size = get_chunked_dim_size(sharding_dim_size, split_size, idx) + if rank == placement.rank(): + start_pos = current_offsets + break + current_offsets += chunk_size + return start_pos, chunk_size + + +def clone_module_parameter(module, param_name): + """ + Clone a parameter from a given existing module. + + Args: + module (:class:`torch.nn.Module`): Module whose parameter needs to be cloned. + param_name (str): Name of the parameter of ``module`` that needs to be cloned. + + Returns: cloned tensor as :class:`torch.nn.Parameter`. + """ + tensor = getattr(module, param_name) + return torch.nn.Parameter(tensor.detach().clone()) + +def gen_binary_op_func(python_op, inplace=False): + src_lines = ['def f(lhs, rhs):'] + if "torch" in python_op: + src_lines.append(f' return {python_op}(lhs, rhs)\n') + elif inplace: + src_lines.append(f' lhs {python_op}= rhs\n return lhs\n') + else: + src_lines.append(f' return lhs {python_op} rhs\n') + + code_str = '\n'.join(src_lines) + g = {'torch': torch} + builtins.exec(code_str, g) + return g["f"] diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_st_common.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_st_common.py new file mode 100644 index 0000000000000000000000000000000000000000..b1e7a23b6f52d9ae34fb50f70e117c23ae1c81bf --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_st_common.py @@ -0,0 +1,66 @@ +# mypy: allow-untyped-defs + +import copy +import random +import torch +from torch.distributed._shard import sharded_tensor + +from torch.distributed._shard.sharding_spec import ( + ChunkShardingSpec, +) + +PLACEMENTS = [ + "rank:0/cuda:0", + "rank:1/cuda:1", + "rank:2/cuda:2", + "rank:3/cuda:3", +] + +DEFAULT_GPU_NUM = 4 + + +def _chunk_sharding_specs_list_for_test(sharding_dims, seed=0): + spec_list = [] + for i in range(len(sharding_dims)): + random.Random(seed + i).shuffle(PLACEMENTS) + spec_list.append( + ChunkShardingSpec( + dim=sharding_dims[i], + placements=copy.deepcopy(PLACEMENTS), + ) + ) + return spec_list + +class MyShardedModel2(torch.nn.Module): + def __init__( + self, + spec=None, + group=None, + init_rrefs=True + ) -> None: + super().__init__() + if spec is not None: + self.sharded_tensor2 = sharded_tensor.rand( + spec, 10, 20, process_group=group, init_rrefs=init_rrefs + ) + else: + self.sharded_tensor2 = None + self.random_tensor2 = torch.nn.Parameter(torch.rand(2, 2)) + + +class MyShardedModel1(torch.nn.Module): + def __init__( + self, + spec=None, + group=None, + init_rrefs=True + ) -> None: + super().__init__() + if spec is not None: + self.sharded_tensor1 = sharded_tensor.rand( + spec, 10, 20, process_group=group, init_rrefs=init_rrefs + ) + else: + self.sharded_tensor1 = None + self.random_tensor1 = torch.nn.Parameter(torch.rand(2, 2)) + self.submodule = MyShardedModel2(spec, group, init_rrefs) diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/__init__.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e607665f96ed9f003cdf39e58d36b9580ce40db Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/__pycache__/common_dtensor.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/__pycache__/common_dtensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24eaa57a5d742b6bd5095180765c75b4635ed1a9 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/__pycache__/common_dtensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/common_dtensor.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/common_dtensor.py new file mode 100644 index 0000000000000000000000000000000000000000..64947af2c9cacc68282a1548c03df144ce8c9eea --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/_tensor/common_dtensor.py @@ -0,0 +1,567 @@ +# mypy: allow-untyped-defs + +# Copyright (c) Meta Platforms, Inc. and affiliates + +import itertools +import sys +from dataclasses import dataclass +from functools import partial, wraps +from typing import ( + Any, + Callable, + cast, + Dict, + Iterator, + List, + Sequence, + Tuple, + TypeVar, + Union, +) + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F + +from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate, Shard +from torch.distributed._tensor.placement_types import Placement +from torch.distributed.tensor.parallel import ( + ColwiseParallel, + parallelize_module, + PrepareModuleInput, + RowwiseParallel, + SequenceParallel, +) +from torch.testing._internal.common_distributed import ( + MultiProcessTestCase, + MultiThreadedTestCase, + skip_if_lt_x_gpu, + run_subtests, + TEST_SKIPS, +) + +from torch.utils._pytree import tree_flatten, tree_unflatten, TreeSpec + +DEVICE_TYPE = ( + "cuda" if torch.cuda.is_available() and torch.cuda.device_count() > 1 else "cpu" +) + +NUM_DEVICES = 4 + +# We use this as a proxy for "multiple GPUs exist" +if torch.cuda.is_available() and torch.cuda.device_count() > 1: + # when we actually have multiple GPUs, relax the requirement to smaller counts. + NUM_DEVICES = min(NUM_DEVICES, torch.cuda.device_count()) + +T = TypeVar("T") + + +# simple RMSNorm layer for testing +class RMSNormPython(torch.nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.ones(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x) + return output * self.weight + + +class MLPModule(nn.Module): + def __init__(self, device, bias: bool = True): + super().__init__() + torch.manual_seed(5) + self.net1 = nn.Linear(10, 16, bias=bias, device=device) + self.relu = nn.ReLU() + self.net2 = nn.Linear(16, 10, bias=bias, device=device) + + def forward(self, x): + return self.net2(self.relu(self.net1(x))) + + def reset_parameters(self): + self.net1.reset_parameters() + self.net2.reset_parameters() + + +class MLPStacked(nn.Module): + def __init__(self, device, n_layers: int = 2): + super().__init__() + self.layers = nn.ModuleList([MLPModule(device) for i in range(n_layers)]) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +@dataclass +class ModelArgs: + n_layers: int = 2 + vocab_size: int = 8 + max_seq_len: int = 16 + dim: int = 16 + n_heads: int = 4 + dropout_p: float = 0.1 + use_attn_mask: bool = True + weight_tying: bool = True + checkpoint_activations: bool = False + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + assert args.dim % args.n_heads == 0 + self.head_dim = args.dim // args.n_heads + self.n_heads = args.n_heads + self.dropout_p = args.dropout_p + self.resid_dropout = nn.Dropout(args.dropout_p) + self.use_attn_mask = args.use_attn_mask + + self.wq = nn.Linear(args.dim, args.dim, bias=False) + self.wk = nn.Linear(args.dim, args.dim, bias=False) + self.wv = nn.Linear(args.dim, args.dim, bias=False) + self.wo = nn.Linear(args.dim, args.dim, bias=False) + + def forward(self, x): + bsz, seq_len, _ = x.size() + queries, keys, values = self.wq(x), self.wk(x), self.wv(x) + queries = queries.view(bsz, seq_len, self.n_heads, self.head_dim) + keys = keys.view(bsz, seq_len, self.n_heads, self.head_dim) + values = values.view(bsz, seq_len, self.n_heads, self.head_dim) + + queries = queries.transpose(1, 2) # (bsz, n_heads, seq_len, head_dim) + keys = keys.transpose(1, 2) # (bsz, n_heads, seq_len, head_dim) + values = values.transpose(1, 2) # (bsz, n_heads, seq_len, head_dim) + + output = F.scaled_dot_product_attention( + queries, + keys, + values, + None, + self.dropout_p if self.training else 0, + self.use_attn_mask, + ) + output = output.transpose(1, 2).contiguous().view(bsz, seq_len, -1) + return self.resid_dropout(self.wo(output)) + + +class FeedForward(nn.Module): + def __init__(self, dim, hidden_dim, dropout_p): + super().__init__() + self.w1 = nn.Linear(dim, hidden_dim) + self.gelu = nn.GELU() + self.w2 = nn.Linear(hidden_dim, dim) + self.resid_dropout = nn.Dropout(dropout_p) + + def forward(self, x): + return self.resid_dropout(self.w2(self.gelu(self.w1(x)))) + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.attention_norm = nn.LayerNorm(args.dim) + self.attention = Attention(args) + self.ffn_norm = nn.LayerNorm(args.dim) + self.feed_forward = FeedForward( + args.dim, hidden_dim=4 * args.dim, dropout_p=args.dropout_p + ) + + def forward(self, x): + h = x + self.attention(self.attention_norm(x)) + out = h + self.feed_forward(self.ffn_norm(h)) + return out + + +# A toy transformer model, partly inspired by the nanoGPT model: +# https://github.com/karpathy/nanoGPT. +class Transformer(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + assert args.vocab_size is not None + assert args.max_seq_len is not None + self.model_args = args + self.max_seq_len = args.max_seq_len + self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim) + self.pos_embeddings = nn.Embedding(args.max_seq_len, args.dim) + self.dropout = nn.Dropout(args.dropout_p) + self.layers = nn.ModuleList() + for _ in range(args.n_layers): + self.layers.append(TransformerBlock(args)) + self.norm = nn.LayerNorm(args.dim) + self.output = nn.Linear(args.dim, args.vocab_size, bias=False) + if args.weight_tying: + self.output.weight = self.tok_embeddings.weight + self.checkpoint_activations = args.checkpoint_activations + + def forward(self, tokens): + _bsz, seq_len = tokens.size() + assert seq_len <= self.max_seq_len + h = self.tok_embeddings(tokens) + pos = torch.arange(0, seq_len, device=tokens.device) + p = self.pos_embeddings(pos) # positional embeddings of shape (seq_len, dim) + h = h + p + h = self.dropout(h) + for layer in self.layers: + if self.checkpoint_activations: + h = torch.utils.checkpoint.checkpoint(layer, h, use_reentrant=False) + else: + h = layer(h) + h = self.norm(h) + output = self.output(h).float() + return output + + @staticmethod + def parallelize( + module: "Transformer", device_mesh: DeviceMesh, use_seq_parallel: bool, local_output_for_attn: bool = False + ) -> nn.Module: + assert isinstance(module, Transformer), f"Requires Transformer but got {module}" + # Parallelize the root submodules. + if use_seq_parallel: + root_plan = { + "tok_embeddings": RowwiseParallel(input_layouts=Replicate(), output_layouts=Shard(1)), + "pos_embeddings": RowwiseParallel(input_layouts=Replicate(), output_layouts=Shard(0)), + "norm": SequenceParallel(), + } + else: + root_plan = { + "tok_embeddings": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "pos_embeddings": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + } + + module_tp = parallelize_module(module, device_mesh, root_plan) + # Parallelize the attention and feed forward submodules. + for layer in module_tp.layers: + layer_parallelize_plan = {} + if use_seq_parallel: + layer_parallelize_plan["attention"] = PrepareModuleInput( + input_layouts=Shard(1), + desired_input_layouts=Replicate(), + ) + # shard the RMSNorms + layer_parallelize_plan["attention_norm"] = SequenceParallel() + layer_parallelize_plan["ffn_norm"] = SequenceParallel() + layer_parallelize_plan["attention.wq"] = ColwiseParallel(use_local_output=local_output_for_attn) + layer_parallelize_plan["attention.wk"] = ColwiseParallel(use_local_output=local_output_for_attn) + layer_parallelize_plan["attention.wv"] = ColwiseParallel(use_local_output=local_output_for_attn) + layer_parallelize_plan["attention.wo"] = ( + RowwiseParallel(output_layouts=Shard(1)) + if use_seq_parallel + else RowwiseParallel() + ) + + layer_parallelize_plan["feed_forward.w1"] = ( + ColwiseParallel(input_layouts=Shard(1)) + if use_seq_parallel + else ColwiseParallel() + ) + layer_parallelize_plan["feed_forward.w2"] = ( + RowwiseParallel(output_layouts=Shard(1)) + if use_seq_parallel + else RowwiseParallel() + ) + + parallelize_module(layer, device_mesh, layer_parallelize_plan) + + # Parallelize the output submodule. If weight tying is enabled, we need to + # make sure output.weight is sharded consistently as tok_embeddings.weight, + # at the cost of the all_reduce operation using RowwiseParallel. + output_parallelize_plan = ( + ColwiseParallel( + input_layouts=Shard(1), + output_layouts=Replicate(), + ) + if use_seq_parallel + else ColwiseParallel(output_layouts=Replicate()) + ) + parallelize_module(module_tp.output, device_mesh, output_parallelize_plan) + + if local_output_for_attn: + for layer in module_tp.layers: + layer.attention.n_heads = module_tp.model_args.n_heads // device_mesh.size() + + # Manually set output.weight so that parameters and gradients are shared. + if module_tp.model_args.weight_tying: + module_tp.output.weight = module_tp.tok_embeddings.weight + + return module_tp + + +def skip_unless_torch_gpu(method: T) -> T: + """ + Test decorator which skips the test unless there's a GPU available to torch. + + >>> # xdoctest: +SKIP + >>> @skip_unless_torch_gpu + >>> def test_some_method(self) -> None: + >>> ... + """ + # The builtin @skip_if_no_gpu relies on os.environ['WORLD_SIZE'] being set. + return cast(T, skip_if_lt_x_gpu(NUM_DEVICES)(method)) + + +class DTensorTestBase(MultiProcessTestCase): + @property + def world_size(self) -> int: + return NUM_DEVICES + + @property + def backend(self) -> str: + backend = "nccl" if self.device_type == "cuda" else "gloo" + return backend + + def build_device_mesh(self) -> DeviceMesh: + return DeviceMesh(self.device_type, list(range(self.world_size))) + + def init_pg(self, eager_init) -> None: + if "nccl" in self.backend and torch.cuda.device_count() < self.world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code) + + if self.backend not in ["nccl", "gloo", "mpi", "cpu:gloo,cuda:nccl"]: + raise RuntimeError(f"Backend {self.backend} not supported!") + + device_id = None + if "nccl" in self.backend: + # set device for nccl pg for collectives + torch.cuda.set_device(self.rank) + # we only need to set device_id for nccl backend with eager init + device_id = torch.device(f"{self.device_type}:{self.rank}") if eager_init else None + + # For nccl backend, bind the device to the process if device_id is not None + # so the nccl communicator is immediately formed and we can use `ncclCommSplit` + # for form subgroup to avoid unnecesssary overhead. + dist.init_process_group( + backend=self.backend, + world_size=self.world_size, + rank=self.rank, # pyre-ignore[16] + init_method=f"file://{self.file_name}", # pyre-ignore[16] + device_id=device_id, + ) + + + def destroy_pg(self) -> None: + # Wait for all ranks to reach here before starting shutdown. + # FIXME dist.barrier deadlocks with multiple threads and NCCL: https://github.com/pytorch/pytorch/issues/95895 + # dist.all_reduce(torch.zeros((1,), device="cuda" if torch.cuda.is_available() else "cpu")) + # FIXME can't use the above all_reduce as it causes hangs on bionic and focal. It hangs: + # test_dtensor.py -- DTensorMeshTest.test_dtensor_device_mesh_device_conversion + dist.barrier() + dist.destroy_process_group() + + def setUp(self) -> None: + super().setUp() + self._spawn_processes() + + # pyre-ignore[2]: + def _test_op(self, mesh: DeviceMesh, op_call, *args, **kwargs) -> None: + out = op_call(*args, **kwargs) + dtc = DTensorConverter(mesh, args, kwargs) + for d_args, d_kwargs in dtc: + # pyre can't find assertTrue anymore? + self.assertEqual(dtc.successful(), True) + d_out = op_call(*d_args, **d_kwargs) + self.assertEqual(d_out.full_tensor(), out) + + def run_subtests(self, *args, **kwargs): + return run_subtests(self, *args, **kwargs) + + +TestFunc = Callable[[...], object] + + +# wrapper to initialize comms (processgroup) +def with_comms(eager_init: Union[TestFunc, bool] = False) -> TestFunc: + + def decorator(func, eager_init: bool = False): + + @wraps(func) # pyre-ignore[6] + def wrapper( + self, *args: Tuple[object], **kwargs: Dict[str, Any] # type: ignore[misc] + ) -> None: + # if enough GPU we can use GPU, otherwise we fallback to CPU + if not torch.cuda.is_available() or torch.cuda.device_count() < self.world_size: + self.device_type = "cpu" + else: + self.device_type = DEVICE_TYPE + + self.init_pg(eager_init) + + try: + func(self, *args, **kwargs) # type: ignore[misc] + except Exception as e: + dist.destroy_process_group() + raise e + + self.destroy_pg() + + return wrapper + + return decorator(func=eager_init) if callable(eager_init) else partial(decorator, eager_init=eager_init) + + +class DTensorOpTestBase(MultiThreadedTestCase): + @property + def world_size(self) -> int: + return NUM_DEVICES + + @property + def device_type(self) -> str: + return DEVICE_TYPE + + def build_device_mesh(self): + return DeviceMesh(self.device_type, list(range(self.world_size))) + + def setUp(self) -> None: + super().setUp() + self._spawn_threads() + + +# This is a class for converting args/kwargs of an op into distributed args/kwargs +class DTensorConverter: + def __init__( + self, + mesh: DeviceMesh, + args: Tuple[object, ...], + kwargs: Dict[str, object], + ) -> None: + self.hit = 0 + self.miss = 0 + self.mesh = mesh + self.args = args + self.kwargs = kwargs + flatten_args, flatten_args_spec = tree_flatten(args) + flatten_kwargs, flatten_kwargs_spec = tree_flatten(kwargs) + + self.flatten_args: List[object] = flatten_args + self.flatten_args_spec: TreeSpec = flatten_args_spec + self.flatten_kwargs: List[object] = flatten_kwargs + self.flatten_kwargs_spec: TreeSpec = flatten_kwargs_spec + + choices_for_args = [self.gen_sharding_choices_for_arg(arg) for arg in self.flatten_args if isinstance(arg, torch.Tensor)] + + choices_for_args.extend( + self.gen_sharding_choices_for_arg(arg) for arg in self.flatten_kwargs if isinstance(arg, torch.Tensor) + ) + + self.sharding_combs: Iterator[Sequence[Placement]] = iter( + itertools.product(*choices_for_args) + ) + + def successful(self) -> bool: + return self.hit > 0 and self.miss == 0 + + def is_supported_tensor(self, t: torch.Tensor) -> bool: + # TODO: dist tensor need to support quantized and sparse + # tensors, quantized tensor might be relatively easy, but + # sparse tensor have special layouts that we need to possibly + # deal with, until we are clear about them, we don't officially + # support them. + return not any( + [ + t.is_sparse_csr, + t.is_sparse, + t.is_mkldnn, + t.is_quantized, + t.is_nested, + torch._is_functional_tensor(t), + t.is_neg(), + t.is_conj(), + t.device.type in ("lazy", "meta"), + # We need a way to test if a tensor is batched but there + # is no official APi to do it + # torch._C._is_batched(t), + ] + ) + + def gen_sharding_choices_for_arg(self, arg: torch.Tensor) -> Sequence[Placement]: + mesh_size = self.mesh.size() + sharding_choices: List[Placement] = [Replicate()] + # c10d collective does not support bool tensor + # for bool tensor we treat it as replicated + if arg.dtype != torch.bool: + # only generating choices with: replicate, or sharding + # evenly on a dimension that could be sharded + sharding_choices = sharding_choices + [ + Shard(i) + for i, s in enumerate(arg.shape) + if s > 1 and s % mesh_size == 0 + ] + # TODO: add multi mesh choices + # all_choices = itertools.product( + # *(self.mesh.ndim * [sharding_choices]) + # ) + return sharding_choices + + def __iter__(self) -> "DTensorConverter": + return self + + def __next__(self) -> Tuple[Tuple[object, ...], Dict[str, object]]: + try: + next_sharding_choices = next(self.sharding_combs) + idx = 0 + + new_args: List[object] = [] + for arg in self.flatten_args: + if isinstance(arg, torch.Tensor): + new_args.append( + self.to_dist_tensor( + arg, self.mesh, [next_sharding_choices[idx]] + ) + ) + idx += 1 + else: + new_args.append(arg) + + new_kwargs: List[object] = [] + for arg in self.flatten_kwargs: + if isinstance(arg, torch.Tensor): + new_kwargs.append( + self.to_dist_tensor( + arg, self.mesh, [next_sharding_choices[idx]] + ) + ) + idx += 1 + else: + new_kwargs.append(arg) + + return ( + tree_unflatten(new_args, self.flatten_args_spec), + tree_unflatten(new_kwargs, self.flatten_kwargs_spec), + ) + except StopIteration as e: + raise StopIteration from e + + def to_dist_tensor( + self, t: torch.Tensor, mesh: DeviceMesh, placements: List[Placement] + ) -> torch.Tensor: + if type(t) is torch.Tensor or type(t) is nn.Parameter: + if self.is_supported_tensor(t): + self.hit += 1 + if t.ndim == 0: + # scalar tensor by default will be replicated + r = distribute_tensor(t, mesh, [Replicate()] * mesh.ndim) + else: + # distribute non-scalar tensors + r = distribute_tensor(t, mesh, placements) + if type(t) is nn.Parameter: + r = nn.Parameter( # type: ignore[assignment] + r, requires_grad=r.requires_grad + ) + return r + else: + self.miss += 1 + return t + elif torch.overrides.is_tensor_like(t): + # Blindly converting tensor subclasses to dist tensor can cause + # unpredictable problems, we explicitly disable this conversion + # for now (i.e. we don't support DTensor holding tensor subclass + # until there's a strong reason later). + self.miss += 1 + return t + else: + raise RuntimeError(f"Trying to convert to DTensor, but got {type(t)}") diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/__init__.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59b27012f38f19e438e9a058aba38e0e0e99284a Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/__init__.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db07216905cf0553d9ee5ea53b77d19773925eb6 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/__pycache__/remote_module_test.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/__pycache__/remote_module_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebca62a89cdcda4c80fa93886fb03919fe2d65ef Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/__pycache__/remote_module_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/remote_module_test.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/remote_module_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7a889b0db847351071f6b730437738caef65c350 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/nn/api/remote_module_test.py @@ -0,0 +1,741 @@ +# mypy: allow-untyped-defs + +import enum +from typing import Tuple + +import torch +import torch.distributed.rpc as rpc +import torch.testing._internal.dist_utils as dist_utils +from torch import Tensor, nn +from torch._jit_internal import Future +from torch.distributed.nn import RemoteModule +from torch.distributed.nn.api.remote_module import _REMOTE_MODULE_PICKLED_ATTRIBUTES +from torch.distributed.nn.api.remote_module import _RemoteModule +from torch.testing._internal.common_distributed import skip_if_lt_x_gpu +from torch.testing._internal.common_utils import TemporaryFileName, TEST_WITH_ROCM +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +_PARAM_VAL = torch.nn.Parameter(torch.ones(1)) + + +# RPC handler for querying the device on the destination worker. +def remote_device(module_rref): + for param in module_rref.local_value().parameters(): + return param.device + + +# RPC handler for querying __dict__ on the destination worker. +def remote_module_attributes(remote_module): + return remote_module.__dict__ + + +# RPC handler for running forward on the destination worker. +def remote_forward(remote_module, args): + return remote_module.forward(*args) + +# RPC handler for running forward_async on the destination worker. +def remote_forward_async(remote_module, args): + # Since future cannot be pickled and sent over the RPC layer, + # have to wait and behave just like ``forward_sync``. + return remote_module.forward_async(*args).wait() + +# RPC handler for getting training mode on the destination worker. +def get_remote_training_arg(module_rref): + return module_rref.local_value().training + +class ModuleCreationMode(enum.Enum): + MODULE_CTOR_WITH_INTERFACE = "module_ctor_with_interface" + MODULE_CTOR = "module_ctor" + + +@torch.jit.interface +class MyModuleInterface: + def forward( + self, tensor: Tensor, number: int, word: str = "default" + ) -> Tuple[str, int, Tensor]: + # pyre-ignore[7]: Pyre and torch.jit.interface don't mix well + pass + + +@torch.jit.interface +class RemoteMyModuleInterface: + def forward( + self, tensor: Tensor, number: int, word: str = "default" + ) -> Tuple[str, int, Tensor]: + # pyre-ignore[7]: Pyre and torch.jit.interface don't mix well + pass + + def forward_async( + self, tensor: Tensor, number: int, word: str = "default" + ) -> Future[Tuple[str, int, Tensor]]: + pass + + +class MyModule(nn.Module): + def __init__(self, first_arg, first_kwarg=-1): + super().__init__() + self.param1 = _PARAM_VAL + + def forward( + self, tensor: Tensor, number: int, word: str = "default" + ) -> Tuple[str, int, Tensor]: + return word, number, tensor + + +class BadModule: + def __init__(self, first_arg, first_kwarg=-1): + pass + + +def create_scripted_module(first_arg, first_kwarg=-1): + module = MyModule(first_arg, first_kwarg=first_kwarg) + scripted_module = torch.jit.script(module) + return scripted_module + + +# Common utils for both CPU and CUDA test suites +class CommonRemoteModuleTest(RpcAgentTestFixture): + @property + def world_size(self): # Override setting in RpcAgentTestFixture + return 2 + + @staticmethod + def _create_remote_module_iter(remote_device, modes=None): + if modes is None: + modes = ModuleCreationMode.__members__.values() + + args = (1,) + kwargs = dict(first_kwarg=2) + + if ModuleCreationMode.MODULE_CTOR in modes: + remote_module = RemoteModule(remote_device, MyModule, args, kwargs) + yield remote_module + + if ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE in modes: + remote_module = _RemoteModule( + remote_device, + create_scripted_module, + args, + kwargs, + _module_interface_cls=MyModuleInterface, + ) + scripted_remote_module = torch.jit.script(remote_module) + yield scripted_remote_module + + +class RemoteModuleTest(CommonRemoteModuleTest): + @dist_utils.dist_init + def test_bad_module(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + remote_device = f"{dst_worker_name}/cpu" + args = (1,) + kwargs = dict(first_kwarg=2) + + with self.assertRaisesRegex( + ValueError, + r"Expect `module_cls\(\*args, \*\*kwargs\)` returns an instance of ,", + ): + RemoteModule(remote_device, BadModule, args, kwargs).forward() + + with self.assertRaisesRegex( + ValueError, + r"Expect `module_cls\(\*args, \*\*kwargs\)` returns an instance of ,", + ): + RemoteModule(remote_device, BadModule, args, kwargs).forward() + + + @dist_utils.dist_init + def test_forward_async(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + args = (torch.ones(1), 2, "3") + for remote_module in self._create_remote_module_iter(dst_worker_name): + ret_fut = remote_module.forward_async(*args) + ret = ret_fut.wait() + self.assertEqual(ret, tuple(reversed(args))) + + @dist_utils.dist_init + def test_forward_async_script(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + scripted_remote_module = next( + self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE] + ) + ) + + @torch.jit.script + def run_forward_async(scripted_remote_module: RemoteMyModuleInterface): + ret_fut = scripted_remote_module.forward_async(torch.ones(1), 2, "3") + ret = ret_fut.wait() + return ret + + ret = run_forward_async(scripted_remote_module) + + self.assertEqual(ret, ("3", 2, torch.ones(1))) + + @dist_utils.dist_init + def test_forward_sync(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + args = (torch.ones(1), 2, "3") + for remote_module in self._create_remote_module_iter(dst_worker_name): + ret = remote_module.forward(*args) + self.assertEqual(ret, tuple(reversed(args))) + + @dist_utils.dist_init + def test_forward_sync_script(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + scripted_remote_module = next( + self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE] + ) + ) + + @torch.jit.script + def run_forward(scripted_remote_module: MyModuleInterface): + ret = scripted_remote_module.forward(torch.ones(1), 2, "3") + return ret + + ret = run_forward(scripted_remote_module) + + self.assertEqual(ret, ("3", 2, torch.ones(1))) + + @dist_utils.dist_init + def test_forward_with_kwargs(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + args = (torch.ones(1), 2) + kwargs = dict(word="3") + # Only test Python nn.Module, because script module methods don't support taking kwargs. + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + ret_fut = remote_module.forward_async(*args, **kwargs) + ret = ret_fut.wait() + self.assertEqual(ret, tuple(reversed(args + ("3",)))) + + ret = remote_module.forward(*args, **kwargs) + self.assertEqual(ret, tuple(reversed(args + ("3",)))) + + @dist_utils.dist_init + def test_remote_parameters(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + # Only test Python nn.Module, because script module methods don't support ``remote_parameters``. + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + param_rrefs = remote_module.remote_parameters() + self.assertEqual(len(param_rrefs), 1) + self.assertTrue(torch.equal(param_rrefs[0].to_here(), _PARAM_VAL)) + + @dist_utils.dist_init + def test_get_module_rref(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + # Only test Python nn.Module, because script module methods don't support ``get_module_rref``. + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + rref = remote_module.get_module_rref() + self.assertEqual(rref, remote_module.module_rref) + for param in rref.to_here().parameters(): + self.assertTrue(torch.equal(param, _PARAM_VAL)) + + @dist_utils.dist_init + def test_train_eval(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + remote_module.train() + ret1 = rpc.rpc_sync(dst_worker_name, get_remote_training_arg, args=(remote_module.get_module_rref(),)) + self.assertEqual(ret1, True) + + remote_module.eval() + ret2 = rpc.rpc_sync(dst_worker_name, get_remote_training_arg, args=(remote_module.get_module_rref(),)) + self.assertEqual(ret2, False) + + @dist_utils.dist_init + def test_unsupported_methods(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + with self.assertRaisesRegex( + ValueError, r"Method ``register_buffer`` not supported for RemoteModule" + ): + remote_module.register_buffer("buffer", torch.ones(5)) + with self.assertRaisesRegex( + ValueError, + r"Method ``register_parameter`` not supported for RemoteModule", + ): + remote_module.register_parameter( + "param", torch.nn.Parameter(torch.ones(1)) + ) + with self.assertRaisesRegex( + ValueError, r"Method ``add_module`` not supported for RemoteModule" + ): + remote_module.add_module("empty", None) + + with self.assertRaisesRegex( + ValueError, r"Method ``apply`` not supported for RemoteModule" + ): + fn = torch.rand((3, 3), requires_grad=False) + remote_module.apply(fn) + + with self.assertRaisesRegex( + ValueError, r"Method ``cuda`` not supported for RemoteModule" + ): + remote_module.cuda() + with self.assertRaisesRegex( + ValueError, r"Method ``cpu`` not supported for RemoteModule" + ): + remote_module.cpu() + with self.assertRaisesRegex( + ValueError, r"Method ``type`` not supported for RemoteModule" + ): + remote_module.type(torch.FloatTensor) + with self.assertRaisesRegex( + ValueError, r"Method ``float`` not supported for RemoteModule" + ): + remote_module.float() + with self.assertRaisesRegex( + ValueError, r"Method ``double`` not supported for RemoteModule" + ): + remote_module.double() + with self.assertRaisesRegex( + ValueError, r"Method ``bfloat16`` not supported for RemoteModule" + ): + remote_module.bfloat16() + with self.assertRaisesRegex( + ValueError, r"Method ``to`` not supported for RemoteModule" + ): + remote_module.to("cpu", dtype=torch.int32) + + def hook(module, grad_input, grad_output): + pass + + with self.assertRaisesRegex( + ValueError, + r"Method ``register_backward_hook`` not supported for RemoteModule", + ): + remote_module.register_backward_hook(hook) + with self.assertRaisesRegex( + ValueError, + r"Method ``register_forward_pre_hook`` not supported for RemoteModule", + ): + remote_module.register_forward_pre_hook(hook) + with self.assertRaisesRegex( + ValueError, + r"Method ``register_forward_hook`` not supported for RemoteModule", + ): + remote_module.register_forward_hook(hook) + + with self.assertRaisesRegex( + ValueError, r"Method ``state_dict`` not supported for RemoteModule" + ): + remote_module.state_dict() + with self.assertRaisesRegex( + ValueError, r"Method ``load_state_dict`` not supported for RemoteModule" + ): + remote_module.load_state_dict({}) + + with self.assertRaisesRegex( + ValueError, + r"Method ``parameters`` not supported for RemoteModule. Please use ``remote_parameters`` instead.", + ): + remote_module.parameters() + with self.assertRaisesRegex( + ValueError, + r"Method ``named_parameters`` not supported for RemoteModule", + ): + remote_module.named_parameters() + with self.assertRaisesRegex( + ValueError, r"Method ``buffers`` not supported for RemoteModule" + ): + remote_module.buffers() + with self.assertRaisesRegex( + ValueError, r"Method ``named_buffers`` not supported for RemoteModule" + ): + remote_module.named_buffers() + with self.assertRaisesRegex( + ValueError, r"Method ``children`` not supported for RemoteModule" + ): + remote_module.children() + with self.assertRaisesRegex( + ValueError, r"Method ``named_children`` not supported for RemoteModule" + ): + remote_module.named_children() + with self.assertRaisesRegex( + ValueError, r"Method ``modules`` not supported for RemoteModule" + ): + remote_module.modules() + with self.assertRaisesRegex( + ValueError, r"Method ``named_modules`` not supported for RemoteModule" + ): + remote_module.named_modules() + + with self.assertRaisesRegex( + ValueError, r"Method ``requires_grad_`` not supported for RemoteModule" + ): + remote_module.requires_grad_() + with self.assertRaisesRegex( + ValueError, r"Method ``zero_grad`` not supported for RemoteModule" + ): + remote_module.zero_grad() + with self.assertRaisesRegex( + ValueError, r"Method ``share_memory`` not supported for RemoteModule" + ): + remote_module.share_memory() + with self.assertRaisesRegex( + ValueError, r"Method ``extra_repr`` not supported for RemoteModule" + ): + remote_module.extra_repr() + + @dist_utils.dist_init + def test_send_remote_module_with_a_new_attribute_not_pickled_over_the_wire(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + # If a new attribute is added to this RemoteModule after the initialization, + # and it will be sent over the wire by RPC, + # this new field will not be pickled, because it's not specified in _REMOTE_MODULE_PICKLED_ATTRIBUTES. + # Note that adding a new attribute out of constructor should rarely happen. + # If a new attribute is added to RemoteModule constructor, + # there is a sanity check to enforce developers to add this attribute to either + # _REMOTE_MODULE_PICKLED_ATTRIBUTES or _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING. + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + new_attr_name = "new_attr" + setattr(remote_module, new_attr_name, 1) + + attrs = rpc.rpc_sync( + dst_worker_name, remote_module_attributes, (remote_module,) + ) + self.assertNotIn(new_attr_name, attrs) + + @dist_utils.dist_init + def test_remote_module_py_pickle_not_supported(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + with TemporaryFileName() as fname: + with self.assertRaisesRegex( + RuntimeError, + "Cannot pickle RemoteModule in python pickler. RemoteModule can only be pickled when using RPC", + ): + torch.save(remote_module, fname) + + @dist_utils.dist_init + def test_remote_module_py_pickle_not_supported_script(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE] + ): + with TemporaryFileName() as fname: + with self.assertRaisesRegex(torch.jit.Error, "can only be pickled when using RPC"): + torch.save(remote_module, fname) + + +class ThreeWorkersRemoteModuleTest(CommonRemoteModuleTest): + @property + def world_size(self): # Override setting in CommonRemoteModuleTest + return 3 + + @dist_utils.dist_init + def test_send_remote_module_over_the_wire(self): + if self.rank != 0: + return + dst_worker1_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + dst_worker2_name = dist_utils.worker_name((self.rank + 2) % self.world_size) + + # Unpickled attributes include both the inherent attributes of RemoteModule + # (not inherited from the superclass) and two installed methods. + expected_unpickled_attrs = list(_REMOTE_MODULE_PICKLED_ATTRIBUTES) + expected_unpickled_attrs.append("forward_async") + expected_unpickled_attrs.append("forward") + + # Create a remote module on worker1 and then pass it to worker2 over the RPC layer. + for remote_module in self._create_remote_module_iter( + dst_worker1_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + # Test querying some simple attributes from worker2. + attrs = rpc.rpc_sync( + dst_worker2_name, remote_module_attributes, (remote_module,) + ) + self.assertListEqual(list(attrs.keys()), expected_unpickled_attrs) + self.assertEqual(attrs["on"], "worker1") + self.assertEqual(attrs["device"], "cpu") + self.assertFalse(attrs["is_device_map_set"]) + self.assertFalse(attrs["is_scriptable"]) + + # Test the installed methods on worker1's can be initiated by worker2 over RPC layer. + # NOTE: In practice a remote module should be directly stored on the worker that runs ``forward``` or ``forward_async``, + # not have another worker to initiate forward over the RPC layer. + args = (torch.ones(1), 2, "3") + ret1 = rpc.rpc_sync(dst_worker2_name, remote_forward, (remote_module, args)) + self.assertEqual(ret1, tuple(reversed(args))) + ret2 = rpc.rpc_sync( + dst_worker2_name, remote_forward_async, (remote_module, args) + ) + self.assertEqual(ret2, tuple(reversed(args))) + + @dist_utils.dist_init + def test_send_remote_module_over_the_wire_script_not_supported(self): + if self.rank != 0: + return + dst_worker1_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + dst_worker2_name = dist_utils.worker_name((self.rank + 2) % self.world_size) + + # Unpickled attributes include both the inherent attributes of RemoteModule + # (not inherited from the superclass) and two installed methods. + expected_unpickled_attrs = list(_REMOTE_MODULE_PICKLED_ATTRIBUTES) + expected_unpickled_attrs.append("forward_async") + expected_unpickled_attrs.append("forward") + + with self.assertRaisesRegex( + RuntimeError, "Passing a script RemoteModule over RPC is not supported." + ): + # Create a remote module on worker1 and then pass it to worker2 over the RPC layer. + for remote_module in self._create_remote_module_iter( + dst_worker1_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE] + ): + # Test querying some simple attributes from worker2. + rpc.rpc_sync( + dst_worker2_name, remote_module_attributes, (remote_module,) + ) + + @dist_utils.dist_init + def test_create_remote_module_from_module_rref(self): + if self.rank != 0: + return + dst_worker1_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + dst_worker2_name = dist_utils.worker_name((self.rank + 2) % self.world_size) + + # Create a remote module on worker1 and then pass its `module_rref` to worker2 over the RPC layer. + for remote_module in self._create_remote_module_iter( + dst_worker1_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + remote_module2 = rpc.rpc_sync( + dst_worker2_name, + RemoteModule.init_from_module_rref, + (dst_worker2_name, remote_module.get_module_rref()), + ) + + args = (torch.ones(1), 2, "3") + ret1 = rpc.rpc_sync( + dst_worker1_name, remote_forward, (remote_module, args) + ) + ret2 = rpc.rpc_sync( + dst_worker2_name, remote_forward, (remote_module2, args) + ) + self.assertEqual(ret1, ret2) + + +class CudaRemoteModuleTest(CommonRemoteModuleTest): + @skip_if_lt_x_gpu(1) + @dist_utils.dist_init + def test_valid_device(self): + if self.rank != 0: + return + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = dist_utils.worker_name(dst_rank) + + for remote_module in self._create_remote_module_iter( + f"{dst_worker_name}/cuda:0", modes=[ModuleCreationMode.MODULE_CTOR] + ): + device = rpc.rpc_sync( + dst_worker_name, remote_device, (remote_module.module_rref,) + ) + self.assertEqual(device.type, "cuda") + self.assertEqual(device.index, 0) + + # Test rank works as well. + for remote_module in self._create_remote_module_iter( + f"rank:{dst_rank}/cuda:0", modes=[ModuleCreationMode.MODULE_CTOR] + ): + device = rpc.rpc_sync( + dst_worker_name, remote_device, (remote_module.module_rref,) + ) + self.assertEqual(device.type, "cuda") + self.assertEqual(device.index, 0) + + @skip_if_lt_x_gpu(1) + @dist_utils.dist_init + def test_invalid_devices(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + with self.assertRaisesRegex( + RuntimeError, + r"Expected one of .+ device type at start of device string", + ): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/foo", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + if TEST_WITH_ROCM: + errorString = (r"HIP error: invalid device ordinal\n" + r"HIP kernel errors might be asynchronously reported at some other API call, " + r"so the stacktrace below might be incorrect.\n" + r"For debugging consider passing AMD_SERIALIZE_KERNEL=3") + else: + errorString = r"CUDA error: invalid device ordinal" + with self.assertRaisesRegex( + RuntimeError, errorString + ): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/cuda:100", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex(RuntimeError, r"Invalid device string: 'cpu2'"): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/cpu2", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex(RuntimeError, r"Device string must not be empty"): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex( + ValueError, + r"Could not parse remote_device: worker1/cuda:0/cuda:1. The valid format is '/'", + ): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/cuda:0/cuda:1", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex( + ValueError, + r"Could not parse remote_device: /. The valid format is '/'", + ): + [ + m.forward() + for m in self._create_remote_module_iter( + "/", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex( + ValueError, + r"Could not parse remote_device: /cuda:0. The valid format is '/'", + ): + [ + m.forward() + for m in self._create_remote_module_iter( + "/cuda:0", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + @skip_if_lt_x_gpu(1) + @dist_utils.dist_init + def test_input_moved_to_cuda_device(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + # These two CPU tensors (in args and kwargs) should be implicitly moved to an appropriate cuda device. + t1 = torch.ones(1) + args = (t1, 2) + t2 = t1 * 2 + kwargs = dict(word=t2) + + # Only test Python nn.Module, because script module methods don't support taking kwargs. + for remote_module in self._create_remote_module_iter( + f"{dst_worker_name}/cuda:0", modes=[ModuleCreationMode.MODULE_CTOR] + ): + ret_fut = remote_module.forward_async(*args, **kwargs) + ret = ret_fut.wait() + self.assertEqual(ret, tuple(reversed(args + (t2,)))) + # TODO: Once the RPC backend can support directly sending GPU tensors, the expected device type should be "cuda:0". + self.assertEqual(ret[0].device.type, "cpu") + self.assertEqual(ret[2].device.type, "cpu") + + ret = remote_module.forward(*args, **kwargs) + self.assertEqual(ret, tuple(reversed(args + (t2,)))) + # TODO: Once the RPC backend can support directly sending GPU tensors, the expected device type should be "cuda:0". + self.assertEqual(ret[0].device.type, "cpu") + self.assertEqual(ret[2].device.type, "cpu") + + @skip_if_lt_x_gpu(1) + @dist_utils.dist_init + def test_input_moved_to_cuda_device_script(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + scripted_remote_module = next( + self._create_remote_module_iter( + f"{dst_worker_name}/cuda:0", + modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE], + ) + ) + + @torch.jit.script + def run_forward(scripted_remote_module: MyModuleInterface): + ret = scripted_remote_module.forward(torch.ones(1), 2, "3") + return ret + + ret = run_forward(scripted_remote_module) + + self.assertEqual(ret, ("3", 2, torch.ones(1))) + # TODO: Once the RPC backend can support directly sending GPU tensors, the expected device type should be "cuda:0". + self.assertEqual(ret[2].device.type, "cpu") diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__init__.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f7a3735b57e7bacd3589099bcecd2ab495b6d47 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/dist_optimizer_test.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/dist_optimizer_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee6b1fd49195e096af4c44258606fe246b2c752a Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/dist_optimizer_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/faulty_agent_rpc_test.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/faulty_agent_rpc_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff58aaa1044f9dfa77d6c86c6b7d65cd6f25d0c1 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/faulty_agent_rpc_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/faulty_rpc_agent_test_fixture.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/faulty_rpc_agent_test_fixture.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b2a6618ef1f25ac1fb992ffb2d8b4e1cf16b0896 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/faulty_rpc_agent_test_fixture.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/rpc_agent_test_fixture.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/rpc_agent_test_fixture.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f9d072c9066f3b5a25b2bcb4d61646d28afa809 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/rpc_agent_test_fixture.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/tensorpipe_rpc_agent_test_fixture.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/tensorpipe_rpc_agent_test_fixture.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..494825f77ce7dfd1f15f926c31ffdd17ca54b4c9 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/__pycache__/tensorpipe_rpc_agent_test_fixture.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/dist_autograd_test.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/dist_autograd_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e113b87ccebda81e4709b1151097c229af066cf0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/dist_autograd_test.py @@ -0,0 +1,2777 @@ +# mypy: allow-untyped-defs + +import sys +import threading +import time +from enum import Enum +import random +import torch +import torch.nn as nn +from datetime import timedelta +import torch.distributed as dist +import torch.distributed.autograd as dist_autograd +import torch.distributed.rpc as rpc +import torch.testing._internal.dist_utils +from torch.autograd import Function +from torch.autograd.function import once_differentiable +from torch.distributed.rpc import RRef +from torch.testing._internal.common_utils import IS_MACOS, skip_but_pass_in_sandcastle_if +from torch.testing._internal.dist_utils import ( + dist_init, + initialize_pg, + wait_until_node_failure, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) +from torch.testing._internal.common_distributed import skip_if_lt_x_gpu + + +# Right now we test up to 3-layer nested rpc calls. +# rpc_done[1] and ctx_ids[1] represent rpc is done in prev rank, and context id +# sent from prev rank respectively. +# rpc_done[2] and ctx_ids[2] represents for prev of prev rank. +# rpc_done[3] and ctx_ids[3] represents for prev of prev of prev rank. +# rpc_done[0] and ctx_ids[0] represents for current rank, but mostly not used. +rpc_done = [False, False, False, False] +ctx_ids = [-1, -1, -1, -1] + +known_context_ids = set() + +requires_grad_tensor = torch.ones(3, 3, requires_grad=True) + +# Send rpc done info and context_id to +# dst_rank = (self.rank + rank_distance) % self.world_size +# we don't need a lock here since the GIL is held while executing remote +# python UDFs, so access is serialized across several workers. +def _set_rpc_done(ctx_id, rank_distance): + global rpc_done + global ctx_ids + global known_context_ids + rpc_done[rank_distance] = True + ctx_ids[rank_distance] = ctx_id + known_context_ids.add(ctx_id) + + +def _check_rpc_done(rank_distance): + while not rpc_done[rank_distance]: + time.sleep(0.1) + + +def _torch_ones(sizes, requires_grad=False): + return torch.ones(sizes, requires_grad=requires_grad) + +# This method must be called on the rref owner, and verifies that the grad of +# rref tensor equals to the given grad. +def _compare_owner_value(context_id, rref, grad): + grads = dist_autograd.get_gradients(context_id) + x = grads[rref.local_value()] + if x.is_sparse: + assert grad.is_sparse + x = x.to_dense() + grad = grad.to_dense() + else: + assert not grad.is_sparse + return torch.equal(x, grad) + + +def create_tensor(): + return torch.ones((3, 3), requires_grad=True) + + +def build_sparse_tensor(coalesce=False, requires_grad=True, dtype=torch.float32): + i = [[0, 1, 1], [2, 0, 2]] + v = [3.2, 4.1, 5.3] + tensor = torch.sparse_coo_tensor( + i, v, (3, 3), requires_grad=requires_grad, dtype=dtype + ) + if coalesce: + tensor = tensor.coalesce() + return tensor + + +@torch.jit.script +def create_torchscript_tensor() -> torch.Tensor: + return torch.ones((3, 3)).requires_grad_() + + +def my_py_add(t1, t2): + return torch.add(t1, t2) + + +def my_scalar_add(a, b): + return a + b + + +def my_rref_add(rref_t1, t2): + ret = torch.add(rref_t1.local_value(), t2) + return ret + + +@torch.jit.script +def my_script_add(t1, t2): + return torch.add(t1, t2) + + +@torch.jit.script +def my_script_ref_add(ref_t1: RRef[torch.Tensor], t2: torch.Tensor) -> torch.Tensor: + t1 = ref_t1.to_here() + return torch.add(t1, t2) + + +def my_nested_rref_add(dst, rref_t1, t2): + return rpc.rpc_sync(dst, my_rref_add, args=(rref_t1, t2)) + + +def ret_requires_grad(): + return requires_grad_tensor + + +def my_py_nested_call(t1, t2, dst, world_size, hops): + next_dst = (dst + 1) % world_size + if hops > 0: + return rpc.rpc_sync( + worker_name(next_dst), + my_py_nested_call, + args=(t1, t2, next_dst, world_size, hops - 1), + ) + else: + return rpc.rpc_sync(worker_name(next_dst), my_py_add, args=(t1, t2)) + + +# after dist autograd context is cleaned up, it should be cleaned up on other +# nodes. This helper allows timeout_seconds for those RPCs to be completed, and +# ensures that all the contexts have been cleaned up in that timeframe.any +def _all_contexts_cleaned_up(timeout_seconds=10): + global known_context_ids + start = time.time() + context_id_to_raised = set() + while ( + time.time() - start < timeout_seconds + and context_id_to_raised != known_context_ids + ): + for context_id in known_context_ids: + try: + dist_autograd._retrieve_context(context_id) + except RuntimeError: + context_id_to_raised.add(context_id) + # all contexts have been cleaned up if trying to retrieve any context resulted in a RuntimeError. + success = context_id_to_raised == known_context_ids + return success + + +# This function creates a dis autograd context, run rpc_sync on the given ps, +# and then blocks until the ps has verified the grads are correctly accumulated. +def _run_trainer(rref_t1, t2, ps, rank_diff, sparse): + with dist_autograd.context() as context_id: + ret = rpc.rpc_sync(ps, my_rref_add, args=(rref_t1, t2)) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + dist_autograd.backward(context_id, [loss]) + # prevent deleting dist autograd context + rpc.rpc_sync(ps, _set_rpc_done, args=(context_id, rank_diff)) + rpc.rpc_sync(ps, _check_rpc_done, args=(0,)) + +# This function is the same as _run_trainer, except rpc calls torchscript +# function "my_script_ref_add" instead of python function "my_rref_add" +def _run_trainer_torchscript(rref_t1, t2, ps, rank_diff, sparse): + with dist_autograd.context() as context_id: + ret = rpc.rpc_sync(ps, my_script_ref_add, args=(rref_t1, t2)) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + dist_autograd.backward(context_id, [loss]) + # prevent deleting dist autograd context + rpc.rpc_sync(ps, _set_rpc_done, args=(context_id, rank_diff)) + rpc.rpc_sync(ps, _check_rpc_done, args=(0,)) + + +class SimulateBackwardError(Function): + _simulate_error = True + + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + @once_differentiable + def backward(ctx, input): + if SimulateBackwardError._simulate_error: + raise Exception("Simulate error on backward pass") # noqa: TRY002 + else: + return input + + +class ExecMode(Enum): + LOCAL = 1 # Run the operation locally. + RPC_SYNC = 2 # Run the operation using rpc_sync + REMOTE = 3 # Run the operation using remote. + RPC_ASYNC = 4 # Run the operation using rpc_async + + +# Common utils for both CPU and CUDA test suites +class CommonDistAutogradTest(RpcAgentTestFixture): + def _exec_func_with_dst(self, dst, exec_mode, method, *args): + if ExecMode.LOCAL == exec_mode: + if len(args) == 1 and isinstance(args[0], list): + return method(*args[0]) + return method(*args) + elif ExecMode.RPC_SYNC == exec_mode: + return rpc.rpc_sync(worker_name(dst), method, args=(args)) + elif ExecMode.REMOTE == exec_mode: + return rpc.remote(worker_name(dst), method, args=(args)).to_here() + elif ExecMode.RPC_ASYNC == exec_mode: + fut = rpc.rpc_async(worker_name(dst), method, args=(args)) + return fut.wait() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + def _exec_func(self, exec_mode, method, *args): + return self._exec_func_with_dst( + self._next_rank(), exec_mode, method, *args + ) + + def _next_rank(self): + if hasattr(self, "dst_rank"): + self.dst_rank = (self.dst_rank + 1) % self.world_size + if self.dst_rank == self.rank: + return self._next_rank() + else: + self.dst_rank = (self.rank + 1) % self.world_size + return self.dst_rank + + def _check_rpc_done(self, rank_distance): + _check_rpc_done(rank_distance) + + def _verify_backwards(self, exec_mode, tensors, context_id, local_grads, *args): + if exec_mode == ExecMode.LOCAL: + torch.autograd.backward(tensors) + return [arg.grad for arg in args] + else: + self._verify_backwards_remote(tensors, context_id, local_grads, *args) + + def _verify_backwards_remote(self, tensors, context_id, local_grads, *args): + dist_autograd.backward(context_id, tensors) + + # Verify grads were accumulated appropriately. + grads = dist_autograd.get_gradients(context_id) + nargs = len(args) + ngrads = 0 + for i in range(0, nargs): + if local_grads[i] is not None: + self.assertIn(args[i], grads) + self.assertEqual(local_grads[i], grads[args[i]]) + ngrads += 1 + else: + self.assertNotIn(args[i], grads) + + self.assertEqual(ngrads, len(grads)) + + def _test_graph(self, fn, exec_mode, sparse): + dst_rank = (self.rank + 1) % self.world_size + + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + if sparse: + t1 = build_sparse_tensor() + t2 = build_sparse_tensor() + else: + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync(worker_name(dst_rank), fn, args=(t1, t2)) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + worker_name(dst_rank), fn, args=(t1, t2) + ).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + rpc.rpc_sync( + worker_name(dst_rank), _set_rpc_done, args=(context_id, 1) + ) + + # Verify graph for current context id. + ctx = dist_autograd._current_context() + self.assertEqual(context_id, ctx._context_id()) + send_functions = ctx._send_functions() + self.assertEqual(1, len(send_functions)) + recv_functions = ctx._recv_functions() + self.assertEqual(1, len(recv_functions)) + self._verify_graph_for_first_rpc_call( + next(iter(send_functions.values())), + next(iter(recv_functions.values())), + t1, + t2, + ret, + ) + + # Wait for the prev rank to be done with rpc. + self._check_rpc_done(1) + # Verify graph for previous context id. + ctx = dist_autograd._retrieve_context(ctx_ids[1]) + send_functions = ctx._send_functions() + self.assertEqual(1, len(send_functions)) + self._verify_graph_for_rpc_call_exec(next(iter(send_functions.values()))) + # this barrier is needed so one worker does not clean up their + # autograd context before another worker tries to access it. + dist.barrier() + + # autograd context should be cleaned up by now. + with self.assertRaises(RuntimeError): + ctx = dist_autograd._retrieve_context(context_id) + + # No autograd context available. + with self.assertRaises(RuntimeError): + ctx = dist_autograd._current_context() + + # 3-layer nested calls + def _test_graph_for_py_nested_call(self, exec_mode, sparse): + dst_rank = (self.rank + 1) % self.world_size + + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + if sparse: + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + else: + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + worker_name(dst_rank), + my_py_nested_call, + args=(t1, t2, dst_rank, self.world_size, 1), + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + worker_name(dst_rank), + my_py_nested_call, + args=(t1, t2, dst_rank, self.world_size, 1), + ).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + # Barrier to ensure all RPCs are done. + dist.barrier() + + for rd in [1, 2, 3]: + rpc.rpc_sync( + worker_name((self.rank + rd) % self.world_size), + _set_rpc_done, + args=(context_id, rd), + ) + + # Barrier to ensure all set_rpc_done have completed. + dist.barrier() + + # For self.rank, it has 4 graphs to verify + # One is for current context id when this rank send first rpc call. + # Second one is for prev context id when this rank make 1st nested + # call. + # Third one is for prev prev context id when this rank make + # 2nd nested call. + # Last one is for prev prev prev context id when this rank + # execute the torch.add() operator. + + # Verify first graph for current context id. + ctx = dist_autograd._current_context() + self.assertEqual(context_id, ctx._context_id()) + send_functions = ctx._send_functions() + self.assertEqual(1, len(send_functions)) + recv_functions = ctx._recv_functions() + self.assertEqual(1, len(recv_functions)) + self._verify_graph_for_first_rpc_call( + next(iter(send_functions.values())), + next(iter(recv_functions.values())), + t1, + t2, + ret, + ) + + # Verify second graph for 1st nested call. + ctx = dist_autograd._retrieve_context(ctx_ids[1]) + self._verify_graph_for_nested_rpc_call(ctx) + + # Verify third graph for 2nd nested call. + ctx = dist_autograd._retrieve_context(ctx_ids[2]) + self._verify_graph_for_nested_rpc_call(ctx) + + # verify last graph for rpc call execution. + ctx = dist_autograd._retrieve_context(ctx_ids[3]) + send_functions = ctx._send_functions() + self.assertEqual(1, len(send_functions)) + self._verify_graph_for_rpc_call_exec(next(iter(send_functions.values()))) + # this barrier is needed so one worker does not clean up their + # autograd context before another worker tries to access it. + dist.barrier() + + # Rank0->Rank1->Rank0 + def _test_graph_for_py_nested_call_itself(self, exec_mode, sparse): + dst_rank = (self.rank + 1) % self.world_size + + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + if sparse: + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + else: + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + worker_name(dst_rank), + my_py_nested_call, + args=( + t1, + t2, + (self.rank - 1 + self.world_size) % self.world_size, + self.world_size, + 0, + ), + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + worker_name(dst_rank), + my_py_nested_call, + args=( + t1, + t2, + (self.rank - 1 + self.world_size) % self.world_size, + self.world_size, + 0, + ), + ).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + rpc.rpc_sync( + worker_name((self.rank + 1) % self.world_size), + _set_rpc_done, + args=(context_id, 1), + ) + + # For self.rank, it has 2 graphs to verify. + # One is for current context id when this rank send first rpc + # call and execute the torch.add() operator. + # Another one is for prev context id when this rank make + # nested call. + ctx = dist_autograd._current_context() + self.assertEqual(context_id, ctx._context_id()) + send_functions = ctx._send_functions() + self.assertEqual(2, len(send_functions)) + recv_functions = ctx._recv_functions() + self.assertEqual(2, len(recv_functions)) + self._verify_graph_for_first_rpc_call( + next(iter(send_functions.values())), + list(recv_functions.values())[1], + t1, + t2, + ret, + ) + self._verify_graph_for_rpc_call_exec(list(send_functions.values())[1]) + + # Verify two pairs of send and recv functions for nested + # call + self._check_rpc_done(1) + ctx = dist_autograd._retrieve_context(ctx_ids[1]) + self._verify_graph_for_nested_rpc_call(ctx) + # this barrier is needed so one worker does not clean up their + # autograd context before another worker tries to access it. + dist.barrier() + + def _test_no_graph_with_tensors_not_require_grad(self, exec_mode, sparse): + initialize_pg(self.file_init_method, self.rank, self.world_size) + dst_rank = (self.rank + 1) % self.world_size + with dist_autograd.context() as context_id: + if sparse: + t1 = build_sparse_tensor(requires_grad=False) + t2 = build_sparse_tensor(requires_grad=False) + else: + t1 = torch.ones(3, 3, requires_grad=False) + t2 = torch.zeros(3, 3, requires_grad=False) + if ExecMode.RPC_SYNC == exec_mode: + rpc.rpc_sync( + worker_name(dst_rank), torch.add, args=(t1, t2) + ) + elif ExecMode.REMOTE == exec_mode: + rpc.remote( + worker_name(dst_rank), torch.add, args=(t1, t2) + ).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + rpc.rpc_sync( + worker_name(dst_rank), _set_rpc_done, args=(context_id, 1) + ) + + ctx = dist_autograd._current_context() + send_functions = ctx._send_functions() + self.assertEqual(len(send_functions), 0) + recv_functions = ctx._recv_functions() + self.assertEqual(len(recv_functions), 0) + + # Wait for the prev rank to be done with rpc. + self._check_rpc_done(1) + # NB: RRef.to_here() always passes the autograd context to the + # the callee, as the caller does not know whether the return + # value would contain a requires_grad tensor or not. + # + # rpc/remote with udf (_set_rpc_done here) also always passes the + # autograd context to the callee due to the same reason. + self.assertNotEqual(-1, dist_autograd._retrieve_context(ctx_ids[1])) + dist.barrier() + + def _test_rpc_complex_args(self, exec_mode, sparse): + with dist_autograd.context(): + num_tensors = 10 + tensors = [] + for i in range(num_tensors): + if sparse: + tensor = build_sparse_tensor(requires_grad=(i % 2 == 0)) + else: + tensor = torch.ones(3, 3, requires_grad=(i % 2 == 0)) + tensors.append(tensor) + dst_rank = self._next_rank() + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + worker_name(dst_rank), torch.stack, args=(tensors,) + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + worker_name(dst_rank), torch.stack, args=(tensors,) + ).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + self.assertEqual(torch.stack(tensors), ret) + + # Verify appropriate tensors have been attached the autograd graph. + next_funcs = next(iter(dist_autograd._current_context()._send_functions().values())).next_functions + for i in range(len(next_funcs)): + self.assertEqual( + "torch::autograd::AccumulateGrad", next_funcs[i][0].name() + ) + self.assertEqual(tensors[i], next_funcs[i][0].variable) + + # Verify that the worker id has been recorded in the context + ctx = dist_autograd._current_context() + worker_ids = ctx._known_worker_ids() + self.assertEqual(len(worker_ids), 1) + self.assertEqual(worker_ids, {dst_rank}) + + def context_cleanup_test_helper(self, rpc_args, func, nested=False): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + # test that in dist autograd, in the case that tensors communicated over RPC do + # NOT require grad, we still cleanup the dist autograd contexts created + # on other nodes. This is because the autograd context is still + # communicated over RPC even if tensor arguments do not require grad, as + # it is possible that the response could. + if nested: + dst_rank = (self.rank + 1) % self.world_size + nested_dst_rank = (dst_rank + 1) % self.world_size + dst_ranks = {dst_rank} + else: + dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} + + with dist_autograd.context() as context_id: + for dst_rank in dst_ranks: + rpc.rpc_sync(worker_name(dst_rank), func, args=rpc_args) + rpc.rpc_sync( + worker_name(dst_rank), _set_rpc_done, args=(context_id, 1) + ) + if nested: + rpc.rpc_sync( + worker_name(nested_dst_rank), + _set_rpc_done, + args=(context_id, 2), + ) + # the thread's context id should be cleaned up + with self.assertRaises(RuntimeError): + dist_autograd._retrieve_context(context_id) + # Ensure all peers have finished mutating the + # `known_context_ids` set. + dist.barrier() + # check that all contexts have been cleaned up. + success = _all_contexts_cleaned_up() + self.assertTrue(success) + + def _backward_no_grad_on_tensor(self, t1, t2, sparse): + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + worker_name(self._next_rank()), + torch.add, + args=(t1, t2)) + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + dist_autograd.backward(context_id, [loss], retain_graph=True) + self.assertIsNone(t1.grad) + self.assertIsNone(t2.grad) + + # Now populate .grad with local autograd engine and + # verify dist autograd doesn't mess with it. + loss_local = torch.add(t1, t2) + if sparse: + loss_local = torch.sparse.sum(loss_local) + else: + loss_local = loss_local.sum() + loss_local.backward() + self.assertIsNotNone(t1.grad) + self.assertIsNotNone(t2.grad) + + t1_grad_before = t1.grad + t2_grad_before = t2.grad + dist_autograd.backward(context_id, [loss]) + self.assertEqual(t1_grad_before, t1.grad) + self.assertEqual(t2_grad_before, t2.grad) + + # The current rank first creates a tensor on the rref_owner, and then passes + # the rref with another tensor to the callee to run either my_rref_add or + # my_nested_rref_add, depending on whether the callee is the rref owner. + # The grad of tensor lives on the current rank, and the grad of the rref + # tensor lives on the rref owner. + def _backward_rref(self, callee, rref_owner, t1, t2, local_grads, sparse): + local_ret = torch.add(t1, t2) + if sparse: + local_ret = torch.sparse.sum(local_ret) + else: + local_ret = local_ret.sum() + local_ret.backward() + with dist_autograd.context() as context_id: + if sparse: + rref_t1 = rpc.remote( + rref_owner, build_sparse_tensor, args=(False, True,) + ) + else: + rref_t1 = rpc.remote( + rref_owner, _torch_ones, args=((3, 3),), kwargs={"requires_grad": True} + ) + if callee == rref_owner: + rref = rpc.remote(callee, my_rref_add, args=(rref_t1, t2)) + else: + rref = rpc.remote( + callee, my_nested_rref_add, args=(rref_owner, rref_t1, t2) + ) + ret = rref.to_here() + if sparse: + ret = torch.sparse.sum(ret) + else: + ret = ret.sum() + dist_autograd.backward(context_id, [ret]) + + # verify grads on caller + grads = dist_autograd.get_gradients(context_id) + self.assertIn(t2, grads) + self.assertEqual(grads[t2], t2.grad) + + # verify grads on rref owner + self.assertTrue( + rpc.rpc_sync( + rref_owner, + _compare_owner_value, + args=(context_id, rref_t1, t1.grad), + ) + ) + + # In this test, every rank will serve as a parameter server (ps) and a + # driver, and then kicks off trainers on the other three ranks. So, we have: + # ps = rank0 with trainers = rank1/2/3 + # ps = rank2 with trainers = rank2/3/0 + # ps = rank3 with trainers = rank3/0/1 + # ps = rank4 with trainers = rank0/1/2 + # + # These four test ps-trainer groups run on completely separate autograd + # graphs, but they share the same set of underlying RpcAgents. + def _test_trainer_ps(self, create_ref_fn, trainer_fn, sparse): + if sparse: + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + else: + t1 = torch.ones((3, 3), requires_grad=True) + t2 = torch.zeros((3, 3), requires_grad=True) + + local_ret = torch.add(t1, t2) + if sparse: + torch.sparse.sum(local_ret).backward() + else: + local_ret.sum().backward() + + # create rref on self + rref_t1 = rpc.remote( + worker_name(self.rank), + create_ref_fn, + args=()) + + # kick off forward and backward pass on three other workers (trainers) + rank_diffs = [1, 2, 3] + futures = [ + rpc.rpc_async( + worker_name((self.rank + rank_diff) % self.world_size), + trainer_fn, + args=(rref_t1, t2, worker_name(self.rank), rank_diff, sparse), + ) for rank_diff in rank_diffs + ] + + # check if the trainers have done with their backward pass + for rank_diff in rank_diffs: + self._check_rpc_done(rank_diff) + + # trainers are done and holding the context for verification + for rank_diff in rank_diffs: + # make sure grads are accumulated for the same tensors and values + # are all correct + ctx_id = ctx_ids[rank_diff] + grads = dist_autograd.get_gradients(ctx_id) + local_t1 = rref_t1.to_here() + self.assertIn(local_t1, grads) + self.assertEqual(grads[local_t1], t1.grad) + + # unblock trainers + _set_rpc_done(None, 0) + + # wait until all trainers are done + torch.futures.wait_all(futures) + + def _backward_multiple_round_trips(self, t1, t2, t3, t4, t5, local_grads, sparse): + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + # Multiple RPCs between different nodes. + val = self._exec_func(exec_mode, torch.add, t1, t2) + val = self._exec_func(exec_mode, torch.mul, t3, val) + s1 = self._exec_func(exec_mode, torch.stack, (t4, val)) + s2 = self._exec_func(exec_mode, torch.stack, (t5, val)) + if sparse: + val = self._exec_func(exec_mode, torch.mul, s1, s2) + val = self._exec_func(exec_mode, torch.mul, val, val) + loss = torch.sparse.sum(val) + else: + val = self._exec_func(exec_mode, torch.bmm, s1, s2) + val = self._exec_func(exec_mode, torch.matmul, val, val) + loss = val.sum() + + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2, t3, t4, t5 + ) + local_grads = ret if ret else local_grads + + def _backward_different_dtypes(self, t1, t2, sparse): + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + loss = self._exec_func(exec_mode, torch.add, t1, t2) + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + local_grads = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + + # Run the same code locally and with dist autograd and verify gradients + # are same. + def _backward_simple_python_udf(self, t1, t2, sparse): + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func(exec_mode, my_py_add, t1, t2) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + local_grads = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + + # Run the same code locally and with dist autograd and verify gradients + # are same. + def _backward_simple_script_call(self, t1, t2, sparse): + local_grads = None + for exec_mode in [ + ExecMode.LOCAL, + ExecMode.RPC_SYNC, + ExecMode.RPC_ASYNC, + ExecMode.REMOTE, + ]: + with dist_autograd.context() as context_id: + forward_ret = self._exec_func(exec_mode, my_script_add, t1, t2) + if sparse: + loss = torch.sparse.sum(forward_ret) + else: + loss = forward_ret.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + def _nested_backward_accumulate_grads(self, t1, t2, sparse): + with dist_autograd.context() as context_id: + ret = rpc.rpc_sync( + worker_name(self._next_rank()), + DistAutogradTest._test_nested_backward_accumulate_grads, + args=(t1, t2, self._next_rank()), + ) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + # Run backward twice. + dist_autograd.backward(context_id, [loss], retain_graph=True) + dist_autograd.backward(context_id, [loss]) + + def _backwards_nested_python_udf(self, t1, t2, sparse): + t3 = t1 * t2 + t4 = t1 + t2 + res = t3 + t4 + loss = t1 * t2 * t3 * t4 * res + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + torch.autograd.backward([loss]) + + # Now run distributed autograd. + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + worker_name(self._next_rank()), + DistAutogradTest._nested_python_udf, + args=(t1, t2, self._next_rank()), + ) + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + dist_autograd.backward(context_id, [loss]) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(t1.grad, grads[t1]) + self.assertEqual(t2.grad, grads[t2]) + + def _mixed_requires_grad(self, t1, t2, sparse): + for exec_mode in [ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func( + exec_mode, DistAutogradTest._mixed_requires_grad_operaton, t1, t2 + ) + self.assertEqual(t1 * t2, ret) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + dist_autograd.backward(context_id, [loss]) + self.assertTrue(t1.requires_grad) + self.assertFalse(t2.requires_grad) + grads = dist_autograd.get_gradients(context_id) + self.assertIn(t1, grads) + self.assertNotIn(t2, grads) + self.assertEqual(t2, grads[t1]) + + def _multiple_backward(self, t1, t2, sparse): + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + worker_name(self._next_rank()), + torch.add, + args=(t1, t2)) + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + # Run backward in a loop multiple times. + for _ in range(1000): + dist_autograd.backward(context_id, [loss], retain_graph=True) + + # For current context, this rank sends t1 and t2 tensors to dst_rank, + # then get t3 = torch.add(t1, t2) result tensor. + # For the current context in this rank, it expects graph like this: + # send function: + # rpcSendBackward + # / \ + # t1.AccumulateGrad t2.AccumulateGrad + # + # recv function: + # + # | + # t3.rpcRecvBackward + # + def _verify_graph_for_first_rpc_call( + self, send_function, recv_function, t1, t2, ret + ): + # Retrieve the next functions in the graph. + next_funcs = send_function.next_functions + self.assertEqual(2, len(next_funcs)) + + # We should now hit t1 and t2 in the autograd graph. + self.assertEqual("torch::autograd::AccumulateGrad", next_funcs[0][0].name()) + self.assertEqual(t1, next_funcs[0][0].variable) + self.assertEqual(0, next_funcs[0][1]) + self.assertEqual("torch::autograd::AccumulateGrad", next_funcs[1][0].name()) + self.assertEqual(t2, next_funcs[1][0].variable) + self.assertEqual(0, next_funcs[1][1]) + + # Test recv functions. + self.assertEqual(ret.grad_fn, recv_function) + + # Run the same code locally and with dist autograd and verify gradients + # are same. + def _backward_simple(self, dst, t1, t2, local_grads, sparse): + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func_with_dst( + dst, exec_mode, torch.add, t1, t2 + ) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + # For a context passed from previous nested chain calls, this rank + # receives two tensors t1 and t2, executes torch.add(t1, t2) and sends + # result tensor t3 back. + # For this context in this rank, it expects graph like this: + # send and recv functions: + # rpcSendBackward + # | + # t3.AddBackward0 + # / \ + # t1.recvRpcBackward t2.recvRpcBackward + def _verify_graph_for_rpc_call_exec(self, send_function): + # Verify next function is AddBackward0 + next_funcs = send_function.next_functions + self.assertEqual(1, len(next_funcs)) + add_backward_fn = next_funcs[0][0] + self.assertEqual("AddBackward0", add_backward_fn.name()) + + # Verify the next two functions are the same recv backward function. + next_funcs = add_backward_fn.next_functions + self.assertEqual(2, len(next_funcs)) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[0][0].name() + ) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[1][0].name() + ) + self.assertEqual(next_funcs[0][0], next_funcs[1][0]) + + # For a context passed from previous nested chain calls, this rank + # receives two tensors t1 and t2, forwards t1 and t2 tensors using + # nested rpc call to next dst. In return route, receive result tensor t3 + # from next dst and forwarding t3 back to previous calls. + # For this context in this rank, it expects graph like this: + # send and recv functions for receiving and forwarding t1 and t2: + # rpcSendBackward + # / \ + # t1.recvRpcBackward t2.recvRpcBackward + # send and recv functions for receiving and forwarding t3: + # rpcSendBackward + # | + # t3.recvRpcBackward + def _verify_graph_for_nested_rpc_call(self, ctx): + send_functions = ctx._send_functions() + self.assertEqual(2, len(send_functions)) + + # For send function when making nest rpc call, + # next functions of the send function are two recv functions + # for received two tensors from previous call + next_funcs = next(iter(send_functions.values())).next_functions + self.assertEqual(2, len(next_funcs)) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[0][0].name() + ) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[1][0].name() + ) + self.assertEqual(next_funcs[0][0], next_funcs[1][0]) + + # For send function when returning response to previous call + # next function of the send function is the recv function + # for received tensor result returned from nested call + next_funcs = list(send_functions.values())[1].next_functions + self.assertEqual(1, len(next_funcs)) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[0][0].name() + ) + + +class TensorPipeAgentDistAutogradTest(CommonDistAutogradTest): + + # Sparse tests only work with TensorPipeAgent. + @dist_init + def test_graph_for_builtin_call_sparse(self): + self._test_graph(torch.add, ExecMode.RPC_SYNC, True) + + @dist_init + def test_graph_for_python_call_sparse(self): + self._test_graph(my_py_add, ExecMode.RPC_SYNC, True) + + @dist_init + def test_graph_for_builtin_remote_call_sparse(self): + self._test_graph(torch.add, ExecMode.REMOTE, True) + + @dist_init + def test_graph_for_python_remote_call_sparse(self): + self._test_graph(my_py_add, ExecMode.REMOTE, True) + + @dist_init + def test_graph_for_py_nested_call_sparse(self): + self._test_graph_for_py_nested_call(ExecMode.RPC_SYNC, True) + + @dist_init + def test_graph_for_py_nested_remote_call_sparse(self): + self._test_graph_for_py_nested_call(ExecMode.REMOTE, True) + + @dist_init + def test_graph_for_py_nested_call_itself_sparse(self): + self._test_graph_for_py_nested_call_itself(ExecMode.RPC_SYNC, True) + + @dist_init + def test_graph_for_py_nested_remote_call_itself_sparse(self): + self._test_graph_for_py_nested_call_itself(ExecMode.REMOTE, True) + + @dist_init + def test_no_graph_with_tensors_not_require_grad_sparse(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.RPC_SYNC, True) + + @dist_init + def test_no_graph_with_tensors_not_require_grad_remote_sparse(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.REMOTE, True) + + @dist_init + def test_rpc_complex_args_sparse(self): + self._test_rpc_complex_args(ExecMode.RPC_SYNC, True) + + @dist_init + def test_remote_complex_args_sparse(self): + self._test_rpc_complex_args(ExecMode.REMOTE, True) + + @dist_init + def test_context_cleanup_tensor_with_grad_sparse(self): + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + self.context_cleanup_test_helper(rpc_args=(t1, t2), func=torch.add) + + @dist_init + def test_context_cleanup_tensor_no_grad_sparse(self): + t1 = build_sparse_tensor(requires_grad=False) + self.context_cleanup_test_helper(rpc_args=(t1, t1), func=torch.add) + + @dist_init + def test_context_cleanup_nested_rpc_sparse(self): + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + dst_rank = (self.rank + 1) % self.world_size + args = (t1, t2, dst_rank, self.world_size, 0) + self.context_cleanup_test_helper( + rpc_args=args, func=my_py_nested_call, nested=True + ) + + @dist_init + def test_backward_no_grad_on_tensor_sparse(self): + self._backward_no_grad_on_tensor( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True + ) + + @dist_init + def test_backward_simple_sparse(self): + self._backward_simple( + self._next_rank(), + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True + ) + + @dist_init + def test_backward_simple_self_sparse(self): + self._backward_simple( + self.rank, + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True + ) + + @dist_init + def test_backward_rref_multi_sparse(self): + if self.rank > 0: + callee = "worker0" + rref_owner = callee + self._backward_rref( + callee, + rref_owner, + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True + ) + + @dist_init + def test_backward_rref_sparse(self): + callee = worker_name(self._next_rank()) + rref_owner = callee + self._backward_rref( + callee, + rref_owner, + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True + ) + + @dist_init + def test_backward_rref_nested_sparse(self): + callee = worker_name((self.rank + 1) % self.world_size) + rref_owner = worker_name((self.rank + 2) % self.world_size) + self._backward_rref( + callee, + rref_owner, + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True + ) + + @dist_init + def test_trainer_ps_sparse(self): + self._test_trainer_ps( + build_sparse_tensor, + _run_trainer, + True + ) + + @dist_init + def test_backward_multiple_round_trips_sparse(self): + self._backward_multiple_round_trips( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=False), + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=False), + build_sparse_tensor(requires_grad=True), + None, + True + ) + + @dist_init + def test_backward_different_dtypes_sparse(self): + self._backward_different_dtypes( + build_sparse_tensor(requires_grad=True, dtype=torch.float32), + build_sparse_tensor(requires_grad=True, dtype=torch.float64), + True + ) + + @dist_init + def test_backward_simple_python_udf_sparse(self): + self._backward_simple_python_udf( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True + ) + + @dist_init + def test_backward_simple_script_call_sparse(self): + self._backward_simple_script_call( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True + ) + + @dist_init + def test_nested_backward_accumulate_grads_sparse(self): + self._nested_backward_accumulate_grads( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True + ) + + @dist_init + def test_backwards_nested_python_udf_sparse(self): + # Run equivalent of _nested_python_udf locally. + self._backwards_nested_python_udf( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True + ) + + @dist_init + def test_mixed_requires_grad_sparse(self): + self._mixed_requires_grad( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=False), + True + ) + + @dist_init + def test_multiple_backward_sparse(self): + self._multiple_backward( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True + ) + + @dist_init + def test_embedding_bag_with_no_grad_tensors(self): + dst = self._next_rank() + remote_embedding = rpc.remote( + worker_name(dst), + torch.nn.EmbeddingBag, + args=(16, 16), + kwargs={"mode": "sum", "sparse": True}, + ) + local_embedding = torch.nn.EmbeddingBag(16, 16, mode="sum", sparse=True) + + input = torch.LongTensor([1, 2, 4, 5, 4, 3, 2, 9]) + # requires_grad = True to record send/recv functions + per_sample_weights = torch.rand((8), requires_grad=True) + offsets = torch.LongTensor([0, 4]) + + local_res = local_embedding(input, offsets, per_sample_weights) + + # Run backward twice. + torch.autograd.backward([local_res.sum()], retain_graph=True) + torch.autograd.backward([local_res.sum()]) + local_grad = local_embedding.weight.grad + + with dist_autograd.context() as context_id: + res = rpc.rpc_sync( + worker_name(dst), + DistAutogradTest._call_remote_embedding, + args=(remote_embedding, input, offsets, per_sample_weights), + ) + + # Run backward twice to test accumulation of sparse gradients. + dist_autograd.backward(context_id, [res.sum()], retain_graph=True) + dist_autograd.backward(context_id, [res.sum()]) + + remote_grad = rpc.rpc_sync( + worker_name(dst), + DistAutogradTest._get_grad, + args=(remote_embedding, context_id), + ) + + self.assertEqual(local_grad, remote_grad) + + +class DistAutogradTest(CommonDistAutogradTest): + @dist_init + def test_autograd_context(self): + # Verify max possible id. + max_auto_increment = 281474976710655 + self.assertEqual( + max_auto_increment + (self.worker_id << 48), dist_autograd._get_max_id() + ) + + context_ids = [] + for _ in range(200): + with dist_autograd.context() as context_id: + self.assertEqual( + context_id, + dist_autograd._retrieve_context(context_id)._context_id(), + ) + # First 16 bits should be worker_id. + self.assertEqual(self.worker_id, context_id >> 48) + context_ids.append(context_id) + + for context_id in context_ids: + with self.assertRaisesRegex( + RuntimeError, + f"Could not find autograd context with id: {context_id}", + ): + dist_autograd._retrieve_context(context_id) + + @dist_init + def test_nested_context(self): + with dist_autograd.context(): + # Nested contexts not supported. + with self.assertRaisesRegex( + RuntimeError, "Already have an autograd context id for this thread" + ): + with dist_autograd.context(): + pass + + @dist_init + def test_graph_for_builtin_call(self): + self._test_graph(torch.add, ExecMode.RPC_SYNC, False) + + @dist_init + def test_graph_for_python_call(self): + self._test_graph(my_py_add, ExecMode.RPC_SYNC, False) + + @dist_init + def test_graph_for_builtin_remote_call(self): + self._test_graph(torch.add, ExecMode.REMOTE, False) + + @dist_init + def test_graph_for_python_remote_call(self): + self._test_graph(my_py_add, ExecMode.REMOTE, False) + + @dist_init + def test_graph_for_py_nested_call(self): + self._test_graph_for_py_nested_call(ExecMode.RPC_SYNC, False) + + @dist_init + def test_graph_for_py_nested_remote_call(self): + self._test_graph_for_py_nested_call(ExecMode.REMOTE, False) + + @dist_init + def test_graph_for_py_nested_call_itself(self): + self._test_graph_for_py_nested_call_itself(ExecMode.RPC_SYNC, False) + + @dist_init + def test_graph_for_py_nested_remote_call_itself(self): + self._test_graph_for_py_nested_call_itself(ExecMode.REMOTE, False) + + @dist_init + def test_no_graph_with_tensors_not_require_grad(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.RPC_SYNC, False) + + @dist_init + def test_no_graph_with_tensors_not_require_grad_remote(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.REMOTE, False) + + def _test_grad_only_on_return_value(self, exec_mode): + initialize_pg(self.file_init_method, self.rank, self.world_size) + dst_rank = (self.rank + 1) % self.world_size + with dist_autograd.context() as context_id: + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync(worker_name(dst_rank), ret_requires_grad) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + worker_name(dst_rank), ret_requires_grad + ).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + dist_autograd.backward(context_id, [ret.sum()]) + + rpc.rpc_sync( + worker_name(dst_rank), _set_rpc_done, args=(context_id, 1) + ) + + # Wait for the prev rank to be done with rpc. + self._check_rpc_done(1) + grads = dist_autograd.get_gradients(ctx_ids[1]) + self.assertEqual(1, len(grads)) + self.assertIn(requires_grad_tensor, grads) + self.assertEqual(torch.ones_like(ret), grads[requires_grad_tensor]) + # due to the above get_gradients call, ensure that dist autograd + # contexts aren't cleaned up until all workers exit context managers + dist.barrier() + + @dist_init + def test_grad_only_on_return_value(self): + self._test_grad_only_on_return_value(ExecMode.RPC_SYNC) + + @dist_init + def test_grad_only_on_return_value_remote(self): + self._test_grad_only_on_return_value(ExecMode.REMOTE) + + @dist_init + def test_rpc_complex_args(self): + self._test_rpc_complex_args(ExecMode.RPC_SYNC, False) + + @dist_init + def test_remote_complex_args(self): + self._test_rpc_complex_args(ExecMode.REMOTE, False) + + @dist_init + def test_context_cleanup_tensor_with_grad(self): + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + self.context_cleanup_test_helper(rpc_args=(t1, t2), func=torch.add) + + @dist_init + def test_context_cleanup_tensor_no_grad(self): + t1 = torch.ones(3, 3, requires_grad=False) + self.context_cleanup_test_helper(rpc_args=(t1, t1), func=torch.add) + + @dist_init + def test_context_cleanup_no_tensors(self): + self.context_cleanup_test_helper(rpc_args=(1, 1), func=my_scalar_add) + + @dist_init + def test_context_cleanup_nested_rpc(self): + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + dst_rank = (self.rank + 1) % self.world_size + args = (t1, t2, dst_rank, self.world_size, 0) + self.context_cleanup_test_helper( + rpc_args=args, func=my_py_nested_call, nested=True + ) + + @dist_init + def test_worker_ids_recorded(self): + dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} + with dist_autograd.context() as context_id: + # if no tensors require grad, we should still record worker_ids, as + # the autograd context ID is still passed to other workers. + t1 = torch.ones(3, 3, requires_grad=False) + t2 = torch.zeros(3, 3, requires_grad=False) + for dst_rank in dst_ranks: + rpc.rpc_sync(worker_name(dst_rank), torch.add, args=(t1, t2)) + rpc.rpc_sync( + worker_name(dst_rank), _set_rpc_done, args=(context_id, 1) + ) + # all worker_ids in dst_ranks should be recorded. + ctx = dist_autograd._current_context() + worker_ids = ctx._known_worker_ids() + self.assertEqual(worker_ids, dst_ranks) + + # worker_ids should be recorded when tensors do require grad + t1.requires_grad = True + t2.requires_grad = True + for dst_rank in dst_ranks: + rpc.rpc_sync( + worker_name(dst_rank), torch.add, args=(t1, t2) + ) + rpc.rpc_sync( + worker_name(dst_rank), _set_rpc_done, args=(context_id, 1) + ) + # all worker_ids in dst_ranks should be recorded. + worker_ids = ctx._known_worker_ids() + self.assertEqual(worker_ids, dst_ranks) + + @dist_init + def test_dist_autograd_profiling(self): + with dist_autograd.context() as context_id: + t1 = torch.rand(3, 3, requires_grad=True) + t2 = torch.rand(3, 3, requires_grad=True) + loss = rpc.rpc_sync(worker_name(self._next_rank()), torch.add, args=(t1, t2)).sum() + with torch.autograd.profiler.profile() as p: + dist_autograd.backward(context_id, [loss]) + + function_events = p.function_events + + def get_event(partial_key): + return next(event for event in function_events if partial_key in event.name) + + send_event = get_event("SendRpcBackward") + recv_event = get_event("RecvRpcBackward") + backward_event = get_event("torch::distributed::autograd::backward") + # There should be at least 1 send and recv_events each, corresponding to send/recv functions executed. + self.assertEqual(send_event.count, 1) + self.assertEqual(recv_event.count, 1) + # The CPU total for backward event should be great than send and recv, since + # applying those functions in the backwards pass is a subset of the entire backward pass. + self.assertGreater(backward_event.cpu_time_total, send_event.cpu_time_total) + self.assertGreater(backward_event.cpu_time_total, recv_event.cpu_time_total) + + @dist_init + def test_error_in_context(self): + with dist_autograd.context(): + t1 = torch.rand(3, 3, requires_grad=True) + t2 = torch.rand(6, 6, requires_grad=True) + + with self.assertRaises(RuntimeError): + # This should throw an error since matrix sizes don't match. + rpc.rpc_sync( + worker_name(self._next_rank()), torch.matmul, args=(t1, t2) + ) + + @dist_init + def test_backward_no_grad_on_tensor(self): + self._backward_no_grad_on_tensor( + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + False + ) + + @dist_init + def test_backward_simple(self): + self._backward_simple( + self._next_rank(), + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False + ) + + @dist_init + def test_backward_simple_self(self): + self._backward_simple( + self.rank, + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False + ) + + @dist_init + def test_backward_rref(self): + callee = worker_name(self._next_rank()) + rref_owner = callee + self._backward_rref( + callee, + rref_owner, + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False + ) + + @dist_init + def test_backward_rref_multi(self): + if self.rank > 0: + callee = "worker0" + rref_owner = callee + self._backward_rref( + callee, + rref_owner, + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False + ) + + @dist_init + def test_backward_rref_nested(self): + callee = worker_name((self.rank + 1) % self.world_size) + rref_owner = worker_name((self.rank + 2) % self.world_size) + self._backward_rref( + callee, + rref_owner, + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False + ) + + @dist_init + def test_trainer_ps(self): + self._test_trainer_ps( + create_tensor, + _run_trainer, + False + ) + + @dist_init + def test_trainer_ps_torchscript_functions(self): + # TODO, need more investigation + # there is rref leak when shutting down, suspect it is because + # ref as arg is passed to pybind boundary, and the ref is not garbage + # collected by python when calling shutdown() + import torch.distributed.rpc.api as api + api._ignore_rref_leak = True + + self._test_trainer_ps(create_torchscript_tensor, _run_trainer_torchscript, False) + + @dist_init + def test_backward_multiple_round_trips(self): + self._backward_multiple_round_trips( + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3)), + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3)), + torch.rand((3, 3), requires_grad=True), + None, + False + ) + + @dist_init + def test_backward_different_tensor_dims(self): + local_grads = None + t1 = torch.rand((4, 6), requires_grad=True) + t2 = torch.rand((6, 5)) + t3 = torch.rand((5, 7), requires_grad=True) + t4 = torch.rand((7, 9)) + + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + val = self._exec_func(exec_mode, torch.matmul, t1, t2) + val = self._exec_func(exec_mode, torch.linalg.multi_dot, (val, t3, t4)) + loss = val.sum() + + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2, t2, t3, t4 + ) + local_grads = ret if ret else local_grads + + @dist_init + def test_backward_unused_tensors(self): + local_grads = None + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + t3 = torch.rand((3, 3), requires_grad=True) + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + s = self._exec_func(exec_mode, torch.stack, (t1, t2, t3)) + val = self._exec_func( + exec_mode, + torch.matmul, + torch.narrow(s, 0, 0, 1), + torch.narrow(s, 0, 2, 1), + ) + + loss = val.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2, t3 + ) + local_grads = ret if ret else local_grads + + @dist_init + def test_backward_multiple_output_tensors(self): + local_grads = None + t = torch.rand((10, 2), requires_grad=True) + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + tensor_list = self._exec_func(exec_mode, torch.split, t, 2) + t1 = tensor_list[0] + t2 = tensor_list[2] + t3 = tensor_list[4] + + val = self._exec_func(exec_mode, torch.linalg.multi_dot, (t1, t2, t3)) + + loss = val.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t + ) + local_grads = ret if ret else local_grads + + def _run_test_backward_unused_send_function_in_thread(self): + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + + # We don't use the result of an RPC function, as a result the + # backward pass would hang in the "FAST" mode. + rpc.rpc_sync( + worker_name(self._next_rank()), torch.add, args=(t1, t2) + ) + + val = torch.mul(t1, t2) + + # Run backward, this would hang forever. + dist_autograd.backward(context_id, [val.sum()]) + + @dist_init + def test_backward_unused_send_function(self): + # Run the test in a thread which would never finish. + t = threading.Thread( + target=self._run_test_backward_unused_send_function_in_thread + ) + t.daemon = True + t.start() + t.join(10) # Wait for 10s. + + # Verify thread is still alive (indicating backward hasn't completed yet). + self.assertTrue(t.is_alive()) + + @dist_init + def test_backward_autograd_engine_error(self): + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + # Perform some ops before error simulation. + tmp = (t1 + t2) * (t1 + t2) + t3 = SimulateBackwardError.apply(tmp) + + # Run multiple round trips across different nodes and verify the + # original node receives an error thrown on a node deep in the chain. + val = rpc.rpc_sync( + worker_name(self._next_rank()), torch.add, args=(t2, t3) + ) + val = rpc.rpc_sync( + worker_name(self._next_rank()), torch.mul, args=(val, t2) + ) + val = rpc.rpc_sync( + worker_name(self._next_rank()), torch.matmul, args=(val, t2) + ) + val = rpc.rpc_sync( + worker_name(self._next_rank()), torch.div, args=(val, t2) + ) + + with self.assertRaisesRegex( + RuntimeError, "Error on Node [0-9]+: Simulate error on backward pass" + ): + # Run backwards, and validate we receive an error. + dist_autograd.backward(context_id, [val.sum()]) + + @dist_init(clean_shutdown=False) + @skip_but_pass_in_sandcastle_if( + IS_MACOS, + "Test is flaky on MacOS since libuv error handling is not as robust as TCP", + ) + def test_backward_node_failure(self): + rpc._set_rpc_timeout(5) # 5 seconds + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + res = rpc.rpc_sync( + worker_name(self._next_rank()), torch.add, args=(t1, t2) + ) + + # Wait for all RPCs to be done. + dist.barrier() + + # Kill all odd rank nodes. + if self.rank % 2 == 0: + shutdown_error_regex = self.get_shutdown_error_regex() + # Wait for all other nodes to die. + for rank in range(self.world_size): + if rank % 2 != 0: + wait_until_node_failure(rank, shutdown_error_regex) + + # Shutdown sequence is not very well defined and as a result + # we might see any error given by get_shutdown_error_regex() + with self.assertRaisesRegex(RuntimeError, shutdown_error_regex): + # Run backwards, and validate we receive an error since all + # other nodes are dead. + dist_autograd.backward(context_id, [res.sum()]) + else: + # Exit all other nodes. + pass + + @dist_init + def test_backward_without_context(self): + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + + context_id = 100 # dummy context_id + with self.assertRaisesRegex( + RuntimeError, + f"Could not find autograd context with id: {context_id}", + ): + res = rpc.rpc_sync( + worker_name(self._next_rank()), torch.add, args=(t1, t2) + ) + dist_autograd.backward(context_id, [res.sum()]) + + @dist_init + def test_backward_without_rpc(self): + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + t3 = torch.add(t1, t2) + + dist_autograd.backward(context_id, [t3.sum()]) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(2, len(grads)) + self.assertIn(t1, grads) + self.assertIn(t2, grads) + self.assertEqual(torch.ones(3, 3), grads[t1]) + self.assertEqual(torch.ones(3, 3), grads[t2]) + + @dist_init + def test_backward_invalid_args(self): + with dist_autograd.context() as context_id: + + with self.assertRaisesRegex(TypeError, "incompatible function arguments"): + dist_autograd.backward(context_id, None) + + with self.assertRaisesRegex(TypeError, "incompatible function arguments"): + dist_autograd.backward(None, None) + + with self.assertRaisesRegex( + RuntimeError, "No tensors provided for gradient computation" + ): + dist_autograd.backward(context_id, []) + + with self.assertRaisesRegex(RuntimeError, "requires_grad not set on"): + t = torch.rand(3, 3) + dist_autograd.backward(context_id, [t]) + + with self.assertRaisesRegex( + RuntimeError, "is not a scalar, all roots need to be scalar" + ): + t = torch.rand(3, 3, requires_grad=True) + dist_autograd.backward(context_id, [t]) + + with self.assertRaisesRegex( + RuntimeError, "does not have a valid gradient function" + ): + t = torch.rand(1, requires_grad=True) + dist_autograd.backward(context_id, [t]) + + @dist_init + def test_backward_multiple_roots(self): + local_grads = None + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC]: + with dist_autograd.context() as context_id: + r1 = self._exec_func(exec_mode, torch.add, t1, t2).sum() + r2 = self._exec_func(exec_mode, torch.mul, t1, t2).sum() + r3 = self._exec_func(exec_mode, torch.cos, t1).sum() + r4 = self._exec_func(exec_mode, torch.div, t1, t2).sum() + + local_grads = self._verify_backwards( + exec_mode, [r1, r2, r3, r4], context_id, local_grads, t1, t2 + ) + + @dist_init + def test_backward_different_dtypes(self): + self._backward_different_dtypes( + torch.rand((3, 3), requires_grad=True, dtype=torch.float32), + torch.rand((3, 3), requires_grad=True, dtype=torch.float64), + False + ) + + @dist_init + def test_backward_simple_python_udf(self): + self._backward_simple_python_udf( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False + ) + + @dist_init + def test_backward_simple_script_call(self): + self._backward_simple_script_call( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False + ) + + @staticmethod + def _complex_python_udf(t1, t2): + t3 = torch.nn.functional.linear(t1, t2) + t4 = torch.nn.functional.linear(t2, t3) + t5 = torch.nn.functional.linear(t3, t4) + return torch.linalg.multi_dot([t1, t2, t3, t4, t5]) + + @dist_init + def test_backward_complex_python_udf(self): + # Run the same code locally and with dist autograd and verify gradients + # are same. + local_grads = None + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func( + exec_mode, DistAutogradTest._complex_python_udf, t1, t2 + ) + loss = ret.sum() + local_grads = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + + @staticmethod + def _python_udf_with_backward_error(t1, t2): + t3 = t1 + t2 + t4 = SimulateBackwardError.apply(t3) + return torch.linalg.multi_dot([t1, t2, t3, t4]) + + @staticmethod + def _nested_rpc_call_backward_error(t1, t2, dst): + t1 = t1 * t2 + t2 = t1 + t2 + res = rpc.rpc_sync( + worker_name(dst), + DistAutogradTest._python_udf_with_backward_error, + args=(t1, t2), + ) + return torch.linalg.multi_dot([t1, t2, res]) + + @dist_init + def test_backward_python_udf_error(self): + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + worker_name(self._next_rank()), + DistAutogradTest._nested_rpc_call_backward_error, + args=(t1, t2, self._next_rank()), + ) + with self.assertRaisesRegex( + RuntimeError, "Simulate error on backward pass" + ): + dist_autograd.backward(context_id, [loss.sum()]) + + _backward_done = False + + @dist_init(clean_shutdown=False) + @skip_but_pass_in_sandcastle_if( + IS_MACOS, + "Test is flaky on MacOS since libuv error handling is not as robust as TCP", + ) + def test_backward_node_failure_python_udf(self): + # Set a short timeout to quickly time out failed RPCs. + rpc._set_rpc_timeout(5) # 5 seconds + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + + dst = self._next_rank() + res = rpc.rpc_sync( + worker_name(dst), + my_py_nested_call, + args=(t1, t2, dst, self.world_size, 1), + ) + + dist.barrier() + + # Kill rank 2 (last hop of nested rpc) and verify rank 0 receives an error. + if self.rank == 2: + return + + store = dist.distributed_c10d._get_default_store() + if self.rank == 0: + # Wait for rank 2 to die. + shutdown_error_regex = self.get_shutdown_error_regex() + wait_until_node_failure(2, shutdown_error_regex) + # Shutdown sequence is not very well defined and as a result + # we might see any error given by get_shutdown_error_regex(). + with self.assertRaisesRegex(RuntimeError, shutdown_error_regex): + # Run backwards, and validate we receive an error since rank 2 is dead. + dist_autograd.backward(context_id, [res.sum()]) + + # Mark rank 0 is done in the store, since the RPC framework on + # some nodes might be broken at this point. + store.set('test_backward_node_failure_python_udf_rank0_done', "True") + else: + # Wait for backward to finish on rank 0. + store.wait(['test_backward_node_failure_python_udf_rank0_done'], timedelta(seconds=10)) + + @staticmethod + def _nested_python_udf(t1, t2, dst): + t3 = t1 * t2 + t4 = t1 + t2 + res = rpc.rpc_sync(worker_name(dst), my_py_add, args=(t3, t4)) + return t1 * t2 * t3 * t4 * res + + @dist_init + def test_backwards_nested_python_udf(self): + # Run equivalent of _nested_python_udf locally. + self._backwards_nested_python_udf( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False + ) + + _test_clean_context_backward_context_id = None + + class MyBackwardFunc(Function): + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + @once_differentiable + def backward(ctx, input): + assert DistAutogradTest._test_clean_context_backward_context_id is not None + + # Release the context to simulate error (use barrier before releasing + # context to ensure all nodes execute the backward function). + dist.barrier() + dist_autograd._release_context( + DistAutogradTest._test_clean_context_backward_context_id + ) + + # Verify all contexts are cleaned up. + assert _all_contexts_cleaned_up() + + return input + + @dist_init + def test_clean_context_during_backward(self): + """ + This test simulates the situation where the 'backward' call might throw + an exception locally which would lead to the autograd context being + cleaned up if we're using the context manager. As a result, the autograd + context might be cleaned up while some threads are still using the + autograd context. + + It is fine for the 'backward' call to throw an exception in this test, + but the process should not crash. + """ + initialize_pg(self.file_init_method, self.rank, self.world_size) + + context = dist_autograd._new_context() + context_id = context._context_id() + DistAutogradTest._test_clean_context_backward_context_id = context_id + + # Send the context id to all nodes. + for i in range(0, self.world_size): + if i != self.rank: + rank_distance = (i - self.rank + self.world_size) % self.world_size + rpc.rpc_sync( + worker_name(i), + _set_rpc_done, + args=(context_id, rank_distance), + ) + + dist.barrier() + + # Verify all context ids have been received. + self.assertEqual(self.world_size - 1, len(known_context_ids)) + + t1 = torch.rand((3, 3), requires_grad=True) + for i in range(0, 100): + dst = self._next_rank() + t1 = rpc.rpc_sync(worker_name(dst), torch.add, args=(t1, t1)) + + # Call MyBackwardFunc as the first op of the backward pass to + # ensure we release the context early in the backward pass. + t1 = DistAutogradTest.MyBackwardFunc.apply(t1) + self.assertEqual(100, len(context._send_functions())) + + context_id = 100 # dummy context_id + with self.assertRaisesRegex( + RuntimeError, + f"Could not find autograd context with id: {context_id}", + ): + dist_autograd.backward(context_id, [t1.sum()]) + + # HACK: Killing workers since otherwise the autograd engine gets stuck on + # other nodes. The proper fix would be addressing: + # https://github.com/pytorch/pytorch/issues/27643, which would inform + # other nodes about the failure. + # The autograd engine gets stuck on other nodes since they're waiting to + # receive gradients from the node that received an error (and as a + # result it didn't execute the rest of the graph). + dist.barrier() + rpc.shutdown(graceful=False) + sys.exit(0) + + @classmethod + def _call_remote_embedding(cls, embedding_rref, input, offsets, per_sample_weights): + embedding = embedding_rref.local_value() + return embedding(input, offsets, per_sample_weights) + + @classmethod + def _get_grad(cls, embedding_rref, context_id): + embedding = embedding_rref.local_value() + grad_map = dist_autograd.get_gradients(context_id) + return grad_map[embedding.weight] + + @classmethod + def _mixed_requires_grad_operaton(cls, t1, t2): + if t2.requires_grad: + return t1 - t2 + else: + return t1 * t2 + + @dist_init + def test_mixed_requires_grad(self): + self._mixed_requires_grad( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=False), + False + ) + + class TestDebugInfoFunc(Function): + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + @once_differentiable + def backward(ctx, input): + debug_info = dist_autograd._get_debug_info() + assert debug_info is not None + backward_passes = int(debug_info["num_current_backward_passes"]) + + # Hard to validate exact numbers because of the distributed nature. + # We can't use a barrier() here since that would block the single + # CPU thread available for autograd and can cause deadlocks. + assert backward_passes >= 1 and backward_passes <= 4 + return input + + @dist_init + def test_debug_info(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + i = 0 + res = {} + res[i] = t1 + for rank in range(self.world_size): + if rank != self.rank: + res[i + 1] = rpc.rpc_sync( + worker_name(rank), torch.add, args=(res[i], t2) + ) + i += 1 + + # Call custom function in middle of backward pass to ensure all + # nodes are still waiting on a backward(). + res[i + 1] = DistAutogradTest.TestDebugInfoFunc.apply(res[i]) + i += 1 + + for rank in range(self.world_size): + if rank != self.rank: + res[i + 1] = rpc.rpc_sync( + worker_name(rank), torch.add, args=(res[i], t2) + ) + i += 1 + + dist_autograd.backward(context_id, [res[i].sum()]) + + debug_info = dist_autograd._get_debug_info() + num_autograd_context = int(debug_info["num_autograd_contexts"]) + # Need atleast one context and not more than 4. + self.assertTrue(num_autograd_context >= 1 and num_autograd_context <= 4) + + for rd in range(self.world_size - 1): + rpc.rpc_sync( + worker_name((self.rank + rd + 1) % self.world_size), + _set_rpc_done, + args=(context_id, rd + 1), + ) + + dist.barrier() + + # Validate information + debug_info = dist_autograd._get_debug_info() + assert debug_info is not None + self.assertEqual(0, int(debug_info["num_current_backward_passes"])) + # only have `num_current_backward_passes` and `num_autograd contexts` + self.assertTrue(len(debug_info) == 2) + + self.assertTrue(_all_contexts_cleaned_up()) + + # All contexts should be cleaned up. + debug_info = dist_autograd._get_debug_info() + self.assertEqual(0, int(debug_info["num_autograd_contexts"])) + + @staticmethod + def _workload_thread(): + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + t3 = rpc.rpc_sync("worker0", torch.add, args=(t1, t2)) + t4 = rpc.rpc_sync("worker0", torch.mul, args=(t2, t3)) + t5 = rpc.rpc_sync("worker0", torch.matmul, args=(t3, t4)) + t6 = rpc.rpc_sync("worker0", torch.add, args=(t4, t5)) + + dist_autograd.backward(context_id, [t6.sum()]) + + @dist_init + def test_async_dist_autograd(self): + """ + This test ensures async processing for distributed autograd works + appropriately. This is achieved by spawning multiple threads and + hammering a single node with a lot of backward() calls. + """ + + initialize_pg(self.file_init_method, self.rank, self.world_size) + if self.rank != 0: + # All other ranks schedule work on rank 0. + threads = [] + for _ in range(20): + t = threading.Thread(target=DistAutogradTest._workload_thread) + t.start() + threads.append(t) + + for thread in threads: + thread.join() + + dist.barrier() + + @dist_init + def test_backward_accumulate_grads(self): + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + t3 = torch.matmul(t1, t2) + # Run backward twice. + torch.autograd.backward([t3.sum()], retain_graph=True) + torch.autograd.backward([t3.sum()]) + + t3 = rpc.rpc_sync( + worker_name(self._next_rank()), torch.matmul, args=(t1, t2) + ) + # Run backward twice. + dist_autograd.backward(context_id, [t3.sum()], retain_graph=True) + dist_autograd.backward(context_id, [t3.sum()]) + + # Verify the gradients are same for local and remote execution. + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(2, len(grads)) + self.assertIn(t1, grads) + self.assertIn(t2, grads) + self.assertEqual(t1.grad, grads[t1]) + self.assertEqual(t2.grad, grads[t2]) + + @staticmethod + def _test_nested_backward_accumulate_grads(t1, t2, dst_rank): + return rpc.rpc_sync(worker_name(dst_rank), torch.add, args=(t1, t2)) + + @dist_init + def test_nested_backward_accumulate_grads(self): + self._nested_backward_accumulate_grads( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False + ) + + @dist_init + def test_multiple_backward(self): + self._multiple_backward( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False + ) + + @dist_init(clean_shutdown=False) + def test_multiple_backward_with_errors(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + f'worker{self._next_rank()}', + DistAutogradTest._python_udf_with_backward_error, + args=(t1, t2)).sum() + + try: + # Run backward in a loop multiple times. + for i in range(100): + if i < 50: + with self.assertRaisesRegex(RuntimeError, "Simulate error on backward pass"): + dist_autograd.backward(context_id, [loss], retain_graph=True) + elif i > 50: + # Recovered from error. + dist_autograd.backward(context_id, [loss], retain_graph=True) + else: + dist.barrier() + SimulateBackwardError._simulate_error = False + dist.barrier() + finally: + # Sync before resetting flag. + dist.barrier() + + # Reset the flag. + SimulateBackwardError._simulate_error = True + + @dist_init + def test_backward_verify_hooks(self): + t1 = torch.ones((3, 3), requires_grad=True) + # Double the gradient. + t1.register_hook(lambda grad: grad * 2) + t2 = torch.ones((3, 3), requires_grad=True) + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func(exec_mode, torch.matmul, t1, t2) + loss = ret.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + @dist_init + def test_no_grad_copy(self): + ''' + Similar to test in test_autograd.py. + ''' + # create autograd function that saves grad pointer as class static + class MyFunc(Function): + static_grad_ptr = None + + @staticmethod + def forward(ctx, inp1, inp2): + return inp1 + inp2 + + @staticmethod + def backward(ctx, grad): + MyFunc.static_grad_ptr = grad.data_ptr() + return grad, grad + + class MyFuncSingleGrad(Function): + static_grad_ptr = None + + @staticmethod + def forward(ctx, inp): + return inp + + @staticmethod + def backward(ctx, grad): + MyFuncSingleGrad.static_grad_ptr = grad.data_ptr() + return grad + + class NonContGradFunc(Function): + @staticmethod + def forward(ctx, inp1): + ctx.size = inp1.size() + return torch.tensor([1.]) + + @staticmethod + def backward(ctx, grad): + return torch.ones(1).expand(ctx.size) + + a = torch.randn(5, 6, requires_grad=True) + b = torch.randn(5, 6, requires_grad=True) + # non-contiguous grad should be copied + with dist_autograd.context() as context_id: + dist_autograd.backward(context_id, [NonContGradFunc.apply(MyFunc.apply(a, b))]) + grads = dist_autograd.get_gradients(context_id) + self.assertFalse(grads[a].data_ptr() == MyFunc.static_grad_ptr) + self.assertFalse(grads[b].data_ptr() == MyFunc.static_grad_ptr) + + # test case that should trigger no copy for a + with dist_autograd.context() as context_id: + dist_autograd.backward(context_id, [MyFuncSingleGrad.apply(a)[1][0]]) + grads = dist_autograd.get_gradients(context_id) + p_g = MyFuncSingleGrad.static_grad_ptr + p_a = grads[a].data_ptr() + # Verify there was no clone. + self.assertTrue(p_a == p_g) + + # Test case that should trigger copy for both of a,b. This is + # different in the distributed autograd case since we hold + # a reference to all grads in a vector until all accumulation is done. + with dist_autograd.context() as context_id: + dist_autograd.backward(context_id, [MyFunc.apply(a, b)[1][0]]) + grads = dist_autograd.get_gradients(context_id) + p_g = MyFunc.static_grad_ptr + p_a = grads[a].data_ptr() + p_b = grads[b].data_ptr() + # check a,b uses different grad buffer + self.assertFalse(p_a == p_b) + # both should be copied. + self.assertFalse(grads[a].data_ptr() == MyFunc.static_grad_ptr) + self.assertFalse(grads[b].data_ptr() == MyFunc.static_grad_ptr) + + @dist_init + def test_no_grad_copy_sparse(self): + # create autograd function that saves grad pointer as class static + class MyFunc(Function): + static_grad_ptr = None + + @staticmethod + def forward(ctx, inp): + return inp + + @staticmethod + def backward(ctx, grad): + MyFunc.static_grad_ptr = grad._values().data_ptr() + return grad + + class NonContGradFunc(Function): + static_grad_ptr = None + + @staticmethod + def forward(ctx, inp1, inp2): + return inp1 + inp2 + + @staticmethod + def backward(ctx, grad): + # Create a sparse tensor with non-contiguous indices and values + # and return as grad. + v = torch.rand(1, 3) + i = torch.ones(1, 1, dtype=torch.long) + nv = v.expand(8, 3) + ni = i.expand(1, 8) + ngrad = torch.sparse_coo_tensor(ni, nv, (10, 3), dtype=torch.float32) + NonContGradFunc.static_grad_ptr = ngrad._values().data_ptr() + return ngrad, ngrad + + a = torch.randn(10, 3, requires_grad=True) + b = torch.randn(10, 3, requires_grad=True) + input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9]) + offsets = torch.tensor([0, 4]) + import torch.nn.functional as F + + # test case that should trigger no copy for a. + with dist_autograd.context() as context_id: + emb_matrix = MyFunc.apply(a) + loss = F.embedding_bag(emb_matrix, input, offsets, sparse=True).sum() + dist_autograd.backward(context_id, [loss], retain_graph=True) + grads = dist_autograd.get_gradients(context_id) + p_g = MyFunc.static_grad_ptr + p_a = grads[a]._values().data_ptr() + # check a uses the same buffer + self.assertTrue(p_a == p_g) + + # Run backwards multiple times. + for _ in range(10): + dist_autograd.backward(context_id, [loss], retain_graph=True) + + # non-contiguous indices and value, we should trigger a copy. + with dist_autograd.context() as context_id: + emb_matrix = NonContGradFunc.apply(a, b) + loss = F.embedding_bag(emb_matrix, input, offsets, sparse=True).sum() + dist_autograd.backward(context_id, [loss], retain_graph=True) + grads = dist_autograd.get_gradients(context_id) + p_g = NonContGradFunc.static_grad_ptr + p_a = grads[a]._values().data_ptr() + p_b = grads[b]._values().data_ptr() + # check a,b uses different grad buffer + self.assertFalse(p_a == p_b) + # Verify we cloned both grads. + self.assertFalse(p_a == p_g) + self.assertFalse(p_b == p_g) + + # Run backwards multiple times to verify accumulation. + for _ in range(10): + dist_autograd.backward(context_id, [loss], retain_graph=True) + + @dist_init + def test_grad_copy_sparse_indices_extra_ref(self): + # create autograd function that saves grad pointer as class static + class MyFunc(Function): + static_grad_ptr = None + static_grad_indices_ref = None + static_grad_values_ref = None + + @staticmethod + def forward(ctx, inp): + return inp + + @staticmethod + def backward(ctx, grad): + MyFunc.static_grad_ptr = grad._values().data_ptr() + # indices() and values() return views, so holding onto + # references of them would not increment refcount of indices + # and values inside the sparse tensor. + MyFunc.static_grad_indices_ref = grad._indices() + MyFunc.static_grad_values_ref = grad._values() + return grad + + a = torch.randn(10, 3, requires_grad=True) + input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9]) + offsets = torch.tensor([0, 4]) + import torch.nn.functional as F + + with dist_autograd.context() as context_id: + emb_matrix = MyFunc.apply(a) + loss = F.embedding_bag(emb_matrix, input, offsets, sparse=True).sum() + dist_autograd.backward(context_id, [loss], retain_graph=True) + grads = dist_autograd.get_gradients(context_id) + p_g = MyFunc.static_grad_ptr + p_a = grads[a]._values().data_ptr() + self.assertIsNotNone(MyFunc.static_grad_indices_ref) + self.assertIsNotNone(MyFunc.static_grad_values_ref) + # grad would be stolen, since static_grad_indices_ref and + # static_grad_values_ref are holding onto views and don't bump the + # refcount. + self.assertTrue(p_g == p_a) + + @dist_init + def test_post_hooks(self): + self.hook_called_times = 0 + + def post_hook_add_one(output_grads, input_grads): + self.hook_called_times += 1 + return output_grads + + def post_hook_add_two(output_grads, input_grads): + self.hook_called_times += 2 + return output_grads + + t = torch.rand(10, 10, requires_grad=True) + a = t + t + + # Register post hooks + accumulate_grad_0 = a.grad_fn.next_functions[0][0] + accumulate_grad_0.register_hook(post_hook_add_one) + accumulate_grad_0.register_hook(post_hook_add_two) + + accumulate_grad_1 = a.grad_fn.next_functions[1][0] + accumulate_grad_1.register_hook(post_hook_add_two) + + with dist_autograd.context() as context_id: + loss = a.sum() + dist_autograd.backward(context_id, [loss]) + self.assertEqual(5, self.hook_called_times) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(1, len(grads)) + self.assertTrue(t in grads) + + @staticmethod + def _slow_add(t1, t2): + time.sleep(1) + t3 = t1 + t2 + t3.requires_grad = True + return t3 + + @dist_init + def test_thread_local_context_id(self): + t1 = torch.rand((3, 3)) + t2 = torch.rand((3, 3)) + + t3 = t1 + t2 + t3.requires_grad = True + t3.sum().backward() + + dst = worker_name((self.rank + 1) % self.world_size) + rref = rpc.remote(dst, DistAutogradTest._slow_add, args=(t1, t2)) + + with dist_autograd.context() as context_id: + loss = rref.to_here().sum() + # due to slow add, the continuation of this backward pass will be + # invoked by the previous rpc.remote thread which does not have a + # valid context_id. So, this can test whether we propagate + # thread_local states properly when jumping across threads on the + # server side. + dist_autograd.backward(context_id, [loss]) + self.assertTrue( + rpc.rpc_sync( + dst, + _compare_owner_value, + args=(context_id, rref, t3.grad) + ) + ) + + +class CudaDistAutogradTest(CommonDistAutogradTest): + @skip_if_lt_x_gpu(1) + @dist_init + def test_gpu_simple(self): + t1 = torch.rand(3, 3, requires_grad=True, device="cuda:0") + t2 = torch.rand(3, 3, requires_grad=True, device="cuda:0") + (t1 + t2).sum().backward() + with dist_autograd.context() as context_id: + t3 = t1 + t2 + dist_autograd.backward(context_id, [t3.sum()]) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(2, len(grads)) + self.assertEqual(t1.grad, grads[t1]) + self.assertEqual(t2.grad, grads[t2]) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_gpu_to_cpu_continuation(self): + t1 = torch.rand(3, 3, requires_grad=True, device="cuda:0") + t2 = torch.rand(3, 3, requires_grad=True) + # Run a few iterations. + for _ in range(3): + t1.grad = None + t2.grad = None + # Root is CPU + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC]: + with dist_autograd.context() as context_id: + t3 = self._exec_func(exec_mode, torch.add, t2, t2) + t4 = t3.cuda(0) + t1 + t5 = self._exec_func(exec_mode, torch.add, t4.cpu(), t2) + t6 = t5.cuda(0) + t4 + t7 = self._exec_func(exec_mode, torch.add, t6.cpu(), t5) + # Autograd graph consists of CPU -> GPU -> CPU execution. + ret = self._verify_backwards( + exec_mode, [t7.sum()], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + @skip_if_lt_x_gpu(1) + @dist_init + def test_gpu_to_cpu_continuation_gpu_root(self): + t1 = torch.rand(3, 3, requires_grad=True, device="cuda:0") + t2 = torch.rand(3, 3, requires_grad=True) + # Run a few iterations. + for _ in range(3): + t1.grad = None + t2.grad = None + # Root is CPU + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC]: + with dist_autograd.context() as context_id: + t3 = self._exec_func(exec_mode, torch.add, t2, t2) + t4 = t3.cuda(0) + t1 + t5 = self._exec_func(exec_mode, torch.add, t4.cpu(), t2) + t6 = t5.cuda(0) + t4 + # Autograd graph consists of CPU -> GPU -> CPU execution. + ret = self._verify_backwards( + exec_mode, [t6.sum()], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + +class FaultyAgentDistAutogradTest(RpcAgentTestFixture): + # Reusing a simplified helper function from DistAutogradTest to ensure + # autograd context is successfully cleaned up even when RPCs are failing. + def context_cleanup_test_helper(self, rpc_args, func): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + # test that in dist autograd, in the case that tensors communicated over RPC do + # NOT require grad, we still cleanup the dist autograd contexts created + # on other nodes. This is because the autograd context is still + # communicated over RPC even if tensor arguments do not require grad, as + # it is possible that the response could. + dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} + + with dist_autograd.context() as context_id: + for dst_rank in dst_ranks: + rpc.rpc_sync(worker_name(dst_rank), func, args=rpc_args) + rpc.rpc_sync( + worker_name(dst_rank), _set_rpc_done, args=(context_id, 1) + ) + # the thread's context id should be cleaned up + with self.assertRaises(RuntimeError): + dist_autograd._retrieve_context(context_id) + # Ensure all peers have finished mutating the + # `known_context_ids` set. + dist.barrier() + # check that all contexts have been cleaned up. + success = _all_contexts_cleaned_up() + self.assertTrue(success) + + # no faulty_messages defined so this fails all retryable messages - see + # faulty_rpc_agent_test_fixture.py for the list of retryable messages. + @dist_init + def test_context_cleanup_tensor_with_grad(self): + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + self.context_cleanup_test_helper(rpc_args=(t1, t2), func=torch.add) + + @dist_init + def test_verify_backend_options(self): + self.assertEqual(self.rpc_backend, rpc.backend_registry.BackendType.FAULTY_TENSORPIPE) + self.assertEqual(self.rpc_backend_options.num_worker_threads, 8) + self.assertEqual(self.rpc_backend_options.num_fail_sends, 3) + self.assertEqual(len(self.rpc_backend_options.messages_to_fail), 4) + + +class WrapperModule(nn.Module): + def __init__(self, model, device): + super().__init__() + self.model = model.to(device) + + def forward(self, *args): + return self.model(*args) + + def gradients(self, ctx_id): + grads = dist_autograd.get_gradients(ctx_id) + return [grads[p] for p in self.model.parameters()] + + +class TensorPipeCudaDistAutogradTest(RpcAgentTestFixture): + + @skip_if_lt_x_gpu(4) + def test_device_maps_backward_pass(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + + # The reverse of this device mapping should be used for the backward pass. + options.set_device_map(dst, {self.rank: (self.rank + 1) % self.world_size}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + t1 = torch.rand(10, device=self.rank, requires_grad=True) + t2 = torch.rand(10, device=self.rank, requires_grad=True) + with dist_autograd.context() as context_id: + res = rpc.rpc_sync(dst, torch.add, args=(t1, t2)) + dist_autograd.backward(context_id, [res.sum()]) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(torch.ones(10), grads[t1]) + self.assertEqual(torch.ones(10), grads[t2]) + self.assertEqual(t1.device, grads[t1].device) + self.assertEqual(t2.device, grads[t2].device) + + rpc.shutdown() + + class MyRemoteCompute(torch.nn.Module): + def forward(self, input): + input = input * 2.0 + return input + + class MyLocalCompute(torch.nn.Module): + def __init__(self, next_stage): + super().__init__() + self.next_stage = next_stage + + def forward(self, input): + return self.next_stage.rpc_sync().forward(input) + + @skip_if_lt_x_gpu(4) + def test_dist_autograd_sync_streams(self): + + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + + # The reverse of this device mapping should be used for the backward pass. + options.set_device_map(dst, {self.rank: (self.rank + 1) % self.world_size}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + remote_compute = rpc.remote(dst, TensorPipeCudaDistAutogradTest.MyRemoteCompute) + local_compute = TensorPipeCudaDistAutogradTest.MyLocalCompute(remote_compute) + for _ in range(10): + input = torch.rand([1000, 10000], device=self.rank, requires_grad=True) + # Run local autograd + result = input * 2.0 + r = random.random() + loss = result.sum() * r + loss.backward() + + # Run distributed autograd + with dist_autograd.context() as context_id: + result = local_compute(input) + loss = result.sum() * r + dist_autograd.backward(context_id, [loss]) + + # Compare grads. + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(input.grad, grads[input]) + + rpc.shutdown() + + @skip_if_lt_x_gpu(4) + def test_gradients_synchronizations(self): + options = self.rpc_backend_options + for peer_rank in range(self.world_size): + options.set_device_map(worker_name(peer_rank), {self.rank: peer_rank}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + if self.rank == 0: + # this is master + layers = [nn.Linear(2000, 2000) for _ in range(self.world_size - 1)] + local_layers = [l.to(0) for l in layers] + remote_layers = [ + rpc.remote( + worker_name(rank), + WrapperModule, + args=(layers[rank - 1], rank) + ) for rank in range(1, self.world_size) + ] + + x = torch.randn(5000, 2000).to(0) + # local iteration + local_model = nn.Sequential(*local_layers) + local_model(x).sum().backward() + + # remote iteration + with dist_autograd.context() as context_id: + for remote_layer in remote_layers: + x = remote_layer.rpc_sync().forward(x) + + dist_autograd.backward(context_id, [x.sum()]) + + futs = [] + for remote_layer in remote_layers: + futs.append(remote_layer.rpc_async().gradients(context_id)) + + for i in range(len(futs)): + local_gradients = [p.grad for p in local_layers[i].parameters()] + for g1, g2 in zip(futs[i].wait(), local_gradients): + self.assertEqual(g1, g2) + + rpc.shutdown() diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/dist_optimizer_test.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/dist_optimizer_test.py new file mode 100644 index 0000000000000000000000000000000000000000..25a13003b7a0abee0bc4b295d880fff341913728 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/dist_optimizer_test.py @@ -0,0 +1,279 @@ +# mypy: allow-untyped-defs + + +import threading + +import torch +import torch.distributed.autograd as dist_autograd +import torch.distributed.rpc as rpc +from torch import optim +from torch.distributed.optim import DistributedOptimizer +from torch.testing._internal.dist_utils import dist_init +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +class MyModule: + lock = threading.Lock() + + def __init__(self, requires_grad=True): + # cannot directly use torch.manual_seed(0) as all threads share the same + # default generator. The race from multiple RPC threads could mess up + # the draw order from the default RNG instance, leading to + # non-deterministic behavior. Hence, create a dedicated RNG here. + g_cpu = torch.Generator() + g_cpu.manual_seed(0) + self.w = torch.rand((3, 3), requires_grad=requires_grad, generator=g_cpu) + + def forward(self, t1): + return torch.mm(self.w, t1) + + def get_w(self): + return self.w + + +class FailingOptimizer(optim.Optimizer): + def __init__(self, params): + super().__init__(params, {}) + + def step(self, closure=None): + raise ValueError("Error running optimizer.") + + +class OptimizerFailingOnConstructor(optim.Optimizer): + def __init__(self, params): + super().__init__(params, {}) + raise ValueError("Error creating optimizer.") + + def step(self, closure=None): + raise NotImplementedError + + +def _call_method(method, obj_rref, *args, **kwargs): + return method(obj_rref.local_value(), *args, **kwargs) + + +def remote_method(method, obj_rref, *args, **kwargs): + """ + Call rpc.remote on a method in a remote object. + + Args: + method: the method (for example, Class.method) + obj_rref (RRef): remote reference to the object + args: positional arguments to pass to the method + kwargs: keyword arguments to pass to the method + + Returns a RRef to the remote method call result. + """ + return rpc.remote( + obj_rref.owner(), + _call_method, + args=[method, obj_rref] + list(args), + kwargs=kwargs, + ) + + +def rpc_async_method(method, obj_rref, *args, **kwargs): + """ + Call rpc.rpc_async on a method in a remote object. + + Args: + method: the method (for example, Class.method) + obj_rref (RRef): remote reference to the object + args: positional arguments to pass to the method + kwargs: keyword arguments to pass to the method + + Returns a Future to the method call result. + """ + return rpc.rpc_async( + obj_rref.owner(), + _call_method, + args=[method, obj_rref] + list(args), + kwargs=kwargs, + ) + + +class DistOptimizerTest(RpcAgentTestFixture): + @dist_init() + def test_dist_optim_exception(self): + # distributed version + owner1 = "worker%d" % ((self.rank + 1) % self.world_size) + owner2 = "worker%d" % ((self.rank + 2) % self.world_size) + + remote_module1 = rpc.remote(owner1, MyModule) + remote_module2 = rpc.remote(owner2, MyModule) + remote_param1 = remote_method(MyModule.get_w, remote_module1) + remote_param2 = remote_method(MyModule.get_w, remote_module2) + + dist_optim = DistributedOptimizer( + FailingOptimizer, [remote_param1, remote_param2] + ) + + with dist_autograd.context() as context_id: + g_cpu = torch.Generator() + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = rpc_async_method(MyModule.forward, remote_module1, t2) + output2 = rpc_async_method(MyModule.forward, remote_module2, output1.wait()) + loss = torch.add(output2.wait(), t1).sum() + + dist_autograd.backward(context_id, [loss]) + with self.assertRaisesRegex(Exception, "Error running optimizer"): + dist_optim.step(context_id) + + @dist_init() + def test_dist_optim_exception_on_constructor(self): + # distributed version + owner1 = "worker%d" % ((self.rank + 1) % self.world_size) + owner2 = "worker%d" % ((self.rank + 2) % self.world_size) + + remote_module1 = rpc.remote(owner1, MyModule) + remote_module2 = rpc.remote(owner2, MyModule) + remote_param1 = remote_method(MyModule.get_w, remote_module1) + remote_param2 = remote_method(MyModule.get_w, remote_module2) + + with self.assertRaisesRegex(Exception, "Error creating optimizer."): + DistributedOptimizer( + OptimizerFailingOnConstructor, [remote_param1, remote_param2] + ) + + def _test_dist_optim_base(self, optim_cls, *args, **kwargs): + # local version + module1 = MyModule() + module2 = MyModule() + params = [module1.get_w(), module2.get_w()] + local_optim = optim_cls(params, *args, **kwargs) + + old_w1 = module1.w.detach().clone() + old_w2 = module2.w.detach().clone() + + g_cpu = torch.Generator() + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = module1.forward(t2) + output2 = module2.forward(output1) + loss = torch.add(output2, t1).sum() + + loss.backward() + local_optim.step() + + # distributed version + owner1 = "worker%d" % ((self.rank + 1) % self.world_size) + owner2 = "worker%d" % ((self.rank + 2) % self.world_size) + + remote_module1 = rpc.remote(owner1, MyModule) + remote_module2 = rpc.remote(owner2, MyModule) + remote_param1 = remote_method(MyModule.get_w, remote_module1) + remote_param2 = remote_method(MyModule.get_w, remote_module2) + + # sanity check: local and remote initial weights should match + self.assertEqual(old_w1, remote_param1.to_here()) + self.assertEqual(old_w2, remote_param2.to_here()) + + dist_optim = DistributedOptimizer( + optim_cls, [remote_param1, remote_param2], *args, **kwargs + ) + + with dist_autograd.context() as context_id: + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = rpc_async_method(MyModule.forward, remote_module1, t2) + output2 = rpc_async_method(MyModule.forward, remote_module2, output1.wait()) + loss = torch.add(output2.wait(), t1) + + dist_autograd.backward(context_id, [loss.sum()]) + dist_optim.step(context_id) + + new_w1 = rpc_async_method(MyModule.get_w, remote_module1).wait() + new_w2 = rpc_async_method(MyModule.get_w, remote_module2).wait() + + # ensure optimizer changed weights + self.assertNotEqual(old_w1, new_w1) + self.assertNotEqual(old_w2, new_w2) + # ensure local equals remote + self.assertEqual(new_w1, module1.get_w()) + self.assertEqual(new_w2, module2.get_w()) + + @dist_init() + def test_dist_optim(self): + self._test_dist_optim_base(optim.Adagrad, lr=0.05) + self._test_dist_optim_base(optim.Adam, lr=1e-2, amsgrad=True) + self._test_dist_optim_base(optim.AdamW, lr=0.05, amsgrad=True) + self._test_dist_optim_base(optim.SGD, lr=0.05) + self._test_dist_optim_base(optim.SGD, lr=1e-3, momentum=1, weight_decay=1, nesterov=True) + self._test_dist_optim_base(optim.Adadelta, rho=0.95) + self._test_dist_optim_base(optim.RMSprop, lr=0.05) + self._test_dist_optim_base(optim.Adamax, lr=0.05) + self._test_dist_optim_base(optim.Rprop, lr=0.05) + + def _test_dist_optim_none_grads(self, optim_cls, *args, **kwargs): + # local version + module1 = MyModule() + module2 = MyModule(requires_grad=False) + params = [module1.get_w(), module2.get_w()] + local_optim = optim_cls(params, *args, **kwargs) + + old_w1 = module1.w.detach().clone() + old_w2 = module2.w.detach().clone() + + g_cpu = torch.Generator() + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = module1.forward(t2) + output2 = module2.forward(output1) + loss = torch.add(output2, t1).sum() + + loss.backward() + local_optim.step() + + # distributed version + owner1 = "worker%d" % ((self.rank + 1) % self.world_size) + owner2 = "worker%d" % ((self.rank + 2) % self.world_size) + + remote_module1 = rpc.remote(owner1, MyModule) + remote_module2 = rpc.remote(owner2, MyModule, args=(False,)) + remote_param1 = remote_module1.remote().get_w() + remote_param2 = remote_module2.remote().get_w() + + # sanity check: local and remote initial weights should match + self.assertEqual(old_w1, remote_param1.to_here()) + self.assertEqual(old_w2, remote_param2.to_here()) + + dist_optim = DistributedOptimizer( + optim_cls, [remote_param1, remote_param2], *args, **kwargs + ) + + with dist_autograd.context() as context_id: + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = remote_module1.rpc_async().forward(t2) + output2 = remote_module2.rpc_async().forward(output1.wait()) + loss = torch.add(output2.wait(), t1) + + dist_autograd.backward(context_id, [loss.sum()]) + dist_optim.step(context_id) + + new_w1 = remote_module1.rpc_async().get_w().wait() + new_w2 = remote_module2.rpc_async().get_w().wait() + + # ensure optimizer changed weights for w1 + self.assertNotEqual(old_w1, new_w1) + + # ensure optimizer not changed weights for w2 + self.assertEqual(old_w2, new_w2) + # ensure local equals remote + self.assertEqual(new_w1, module1.get_w()) + self.assertEqual(new_w2, module2.get_w()) + + @dist_init() + def test_dist_optim_none_grads(self): + self._test_dist_optim_none_grads(optim.SGD, lr=0.05) + self._test_dist_optim_none_grads(optim.RMSprop, lr=0.05) + self._test_dist_optim_none_grads(optim.Rprop, lr=0.05) + self._test_dist_optim_none_grads(optim.Adadelta, rho=0.95) diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__init__.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c92c3a42f5d2a49768e3903fbfa1eca04fdb895c Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__pycache__/parameter_server_test.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__pycache__/parameter_server_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d00520d9c40e186492f592017177db371e781e69 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__pycache__/parameter_server_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__pycache__/reinforcement_learning_rpc_test.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__pycache__/reinforcement_learning_rpc_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ae5b93b3cf71915db83edee4b83bcf1aa1aa849 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/__pycache__/reinforcement_learning_rpc_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/parameter_server_test.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/parameter_server_test.py new file mode 100644 index 0000000000000000000000000000000000000000..928f28c19211eb4c0277fa31f746b60df0763d91 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/parameter_server_test.py @@ -0,0 +1,140 @@ +# mypy: allow-untyped-defs + +# If you need to modify this file to make this test pass, please also apply same edits accordingly to +# https://github.com/pytorch/examples/blob/master/distributed/rpc/batch/parameter_server.py +# and https://pytorch.org/tutorials/intermediate/rpc_async_execution.html#batch-updating-parameter-server + +import threading +from datetime import datetime +from time import perf_counter + +import torch +import torch.distributed.rpc as rpc +import torch.nn as nn +from torch import optim + +from torch.testing._internal.dist_utils import ( + dist_init, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import RpcAgentTestFixture + +batch_size = 20 +in_features = 100 +out_features = 30 +num_batches = 4 + + +def timed_log(text): + print(f"{datetime.now().strftime('%H:%M:%S')} {text}") + + +class BatchUpdateParameterServer: + + def __init__(self, batch_update_size): + self.model = nn.Linear(in_features, out_features) + self.lock = threading.Lock() + self.future_model = torch.futures.Future() + self.batch_update_size = batch_update_size + self.curr_update_size = 0 + self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9) + for p in self.model.parameters(): + p.grad = torch.zeros_like(p) + + def get_model(self): + return self.model + + @staticmethod + @rpc.functions.async_execution + def update_and_fetch_model(ps_rref, grads): + self = ps_rref.local_value() + for p, g in zip(self.model.parameters(), grads): + if p.grad is None: + p.grad = g + else: + p.grad += g + with self.lock: + timed_log(f"PS got {self.curr_update_size}/{self.batch_update_size} updates") + self.curr_update_size += 1 + fut = self.future_model + + if self.curr_update_size >= self.batch_update_size: + for p in self.model.parameters(): + p.grad /= self.batch_update_size + self.curr_update_size = 0 + self.optimizer.step() + self.optimizer.zero_grad() + fut.set_result(self.model) + timed_log("PS updated model") + self.future_model = torch.futures.Future() + + return fut + + +class Trainer: + + def __init__(self, ps_rref): + self.ps_rref = ps_rref + self.loss_fn = nn.L1Loss() + + def get_next_batch(self): + for _ in range(num_batches): + inputs = torch.randn(batch_size, in_features) + labels = torch.zeros(batch_size, out_features) + yield inputs, labels + + def train(self): + name = rpc.get_worker_info().name + m = self.ps_rref.rpc_sync().get_model() + for inputs, labels in self.get_next_batch(): + timed_log(f"{name} processing one batch") + self.loss_fn(m(inputs), labels).backward() + timed_log(f"{name} reporting grads") + m = rpc.rpc_sync( + self.ps_rref.owner(), + BatchUpdateParameterServer.update_and_fetch_model, + args=(self.ps_rref, [p.grad for p in m.cpu().parameters()]), + ) + timed_log(f"{name} got updated model") + + +def run_trainer(ps_rref): + trainer = Trainer(ps_rref) + trainer.train() + + +def run_ps(trainers): + timed_log("Start training") + start = perf_counter() + ps_rref = rpc.RRef(BatchUpdateParameterServer(len(trainers))) + futs = [rpc.rpc_async(trainer, run_trainer, args=(ps_rref,)) for trainer in trainers] + + torch.futures.wait_all(futs) + stop = perf_counter() + timed_log("Finish training") + timed_log(f"Time spent training: {stop-start}s") + +class ParameterServerTest(RpcAgentTestFixture): + + @dist_init(setup_rpc=False) + def test_batch_updating_parameter_server(self): + + if self.rank != 0: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + else: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + run_ps([f"{worker_name(r)}" for r in range(1, self.world_size)]) + + rpc.shutdown() diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/reinforcement_learning_rpc_test.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/reinforcement_learning_rpc_test.py new file mode 100644 index 0000000000000000000000000000000000000000..4b3ba7de3f4ac956b68972e688f8f00eede6143d --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/examples/reinforcement_learning_rpc_test.py @@ -0,0 +1,258 @@ +# mypy: allow-untyped-defs + +# If you need to modify this file to make this test pass, please also apply same edits accordingly to +# https://github.com/pytorch/examples/blob/master/distributed/rpc/rl/main.py +# and https://pytorch.org/tutorials/intermediate/rpc_tutorial.html + +import numpy as np + +import torch +import torch.distributed.rpc as rpc +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.distributed.rpc import RRef, rpc_sync, rpc_async, remote +from torch.distributions import Categorical + +from torch.testing._internal.dist_utils import dist_init, worker_name +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import RpcAgentTestFixture + +TOTAL_EPISODE_STEP = 5000 +GAMMA = 0.1 +SEED = 543 + +def _call_method(method, rref, *args, **kwargs): + r""" + a helper function to call a method on the given RRef + """ + return method(rref.local_value(), *args, **kwargs) + + +def _remote_method(method, rref, *args, **kwargs): + r""" + a helper function to run method on the owner of rref and fetch back the + result using RPC + """ + args = [method, rref] + list(args) + return rpc_sync(rref.owner(), _call_method, args=args, kwargs=kwargs) + + +class Policy(nn.Module): + r""" + Borrowing the ``Policy`` class from the Reinforcement Learning example. + Copying the code to make these two examples independent. + See https://github.com/pytorch/examples/tree/master/reinforcement_learning + """ + def __init__(self) -> None: + super().__init__() + self.affine1 = nn.Linear(4, 128) + self.dropout = nn.Dropout(p=0.6) + self.affine2 = nn.Linear(128, 2) + + self.saved_log_probs = [] + self.rewards = [] + + def forward(self, x): + x = self.affine1(x) + x = self.dropout(x) + x = F.relu(x) + action_scores = self.affine2(x) + return F.softmax(action_scores, dim=1) + + +class DummyEnv: + r""" + A dummy environment that implements the required subset of the OpenAI gym + interface. It exists only to avoid a dependency on gym for running the + tests in this file. It is designed to run for a set max number of iterations, + returning random states and rewards at each step. + """ + def __init__(self, state_dim=4, num_iters=10, reward_threshold=475.0): + self.state_dim = state_dim + self.num_iters = num_iters + self.iter = 0 + self.reward_threshold = reward_threshold + + def seed(self, manual_seed): + torch.manual_seed(manual_seed) + + def reset(self): + self.iter = 0 + return torch.randn(self.state_dim) + + def step(self, action): + self.iter += 1 + state = torch.randn(self.state_dim) + reward = torch.rand(1).item() * self.reward_threshold + done = self.iter >= self.num_iters + info = {} + return state, reward, done, info + + +class Observer: + r""" + An observer has exclusive access to its own environment. Each observer + captures the state from its environment, and send the state to the agent to + select an action. Then, the observer applies the action to its environment + and reports the reward to the agent. + """ + def __init__(self) -> None: + self.id = rpc.get_worker_info().id + self.env = DummyEnv() + self.env.seed(SEED) + + def run_episode(self, agent_rref, n_steps): + r""" + Run one episode of n_steps. + Arguments: + agent_rref (RRef): an RRef referencing the agent object. + n_steps (int): number of steps in this episode + """ + state, _ep_reward = self.env.reset(), 0 + for _ in range(n_steps): + # send the state to the agent to get an action + action = _remote_method(Agent.select_action, agent_rref, self.id, state) + + # apply the action to the environment, and get the reward + state, reward, done, _ = self.env.step(action) + + # report the reward to the agent for training purpose + _remote_method(Agent.report_reward, agent_rref, self.id, reward) + + if done: + break + + +class Agent: + def __init__(self, world_size): + self.ob_rrefs = [] + self.agent_rref = RRef(self) + self.rewards = {} + self.saved_log_probs = {} + self.policy = Policy() + self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2) + self.eps = np.finfo(np.float32).eps.item() + self.running_reward = 0 + self.reward_threshold = DummyEnv().reward_threshold + for ob_rank in range(1, world_size): + ob_info = rpc.get_worker_info(worker_name(ob_rank)) + self.ob_rrefs.append(remote(ob_info, Observer)) + self.rewards[ob_info.id] = [] + self.saved_log_probs[ob_info.id] = [] + + def select_action(self, ob_id, state): + r""" + This function is mostly borrowed from the Reinforcement Learning example. + See https://github.com/pytorch/examples/tree/master/reinforcement_learning + The main difference is that instead of keeping all probs in one list, + the agent keeps probs in a dictionary, one key per observer. + + NB: no need to enforce thread-safety here as GIL will serialize + executions. + """ + probs = self.policy(state.unsqueeze(0)) + m = Categorical(probs) + action = m.sample() + self.saved_log_probs[ob_id].append(m.log_prob(action)) + return action.item() + + def report_reward(self, ob_id, reward): + r""" + Observers call this function to report rewards. + """ + self.rewards[ob_id].append(reward) + + def run_episode(self, n_steps=0): + r""" + Run one episode. The agent will tell each observer to run n_steps. + """ + # make async RPC to kick off an episode on all observers + futs = [ + rpc_async( + ob_rref.owner(), + _call_method, + args=(Observer.run_episode, ob_rref, self.agent_rref, n_steps) + ) for ob_rref in self.ob_rrefs + ] + + # wait until all observers have finished this episode + for fut in futs: + fut.wait() + + def finish_episode(self): + r""" + This function is mostly borrowed from the Reinforcement Learning example. + See https://github.com/pytorch/examples/tree/master/reinforcement_learning + The main difference is that it joins all probs and rewards from + different observers into one list, and uses the minimum observer rewards + as the reward of the current episode. + """ + + # joins probs and rewards from different observers into lists + R, probs, rewards = 0, [], [] + for ob_id in self.rewards: + probs.extend(self.saved_log_probs[ob_id]) + rewards.extend(self.rewards[ob_id]) + + # use the minimum observer reward to calculate the running reward + min_reward = min(sum(self.rewards[ob_id]) for ob_id in self.rewards) + self.running_reward = 0.05 * min_reward + (1 - 0.05) * self.running_reward + + # clear saved probs and rewards + for ob_id in self.rewards: + self.rewards[ob_id] = [] + self.saved_log_probs[ob_id] = [] + + policy_loss, returns = [], [] + for r in rewards[::-1]: + R = r + GAMMA * R + returns.insert(0, R) + returns = torch.tensor(returns) + returns = (returns - returns.mean()) / (returns.std() + self.eps) + for log_prob, R in zip(probs, returns): + policy_loss.append(-log_prob * R) + self.optimizer.zero_grad() + policy_loss = torch.cat(policy_loss).sum() + policy_loss.backward() + self.optimizer.step() + return min_reward + + +def run_agent(agent, n_steps): + while True: + agent.run_episode(n_steps=n_steps) + agent.finish_episode() + + if agent.running_reward > agent.reward_threshold: + print(f"Solved! Running reward is now {agent.running_reward}!") + break + + +class ReinforcementLearningRpcTest(RpcAgentTestFixture): + @dist_init(setup_rpc=False) + def test_rl_rpc(self): + if self.rank == 0: + # Rank 0 is the agent. + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + agent = Agent(self.world_size) + run_agent(agent, n_steps=int(TOTAL_EPISODE_STEP / (self.world_size - 1))) + + # Ensure training was run. We don't really care about whether the task was learned, + # since the purpose of the test is to check the API calls. + self.assertGreater(agent.running_reward, 0.0) + else: + # Other ranks are observers that passively wait for instructions from the agent. + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + rpc.shutdown() diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/faulty_agent_rpc_test.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/faulty_agent_rpc_test.py new file mode 100644 index 0000000000000000000000000000000000000000..132e30e5b5cfb41ab9c0a2a0f0f6f960d648c0f9 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/faulty_agent_rpc_test.py @@ -0,0 +1,326 @@ +# mypy: allow-untyped-defs + +import torch +import time +import torch.distributed.rpc as rpc +from torch.distributed.rpc.api import _delete_all_user_and_unforked_owner_rrefs +from torch.testing._internal.dist_utils import ( + dist_init, + wait_until_pending_futures_and_users_flushed, + wait_until_owners_and_forks_on_rank, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + +def my_sleep_func(seconds=1): + time.sleep(seconds) + return torch.mul(torch.tensor(1), torch.tensor(1)) + +@torch.jit.script +def my_script_func(tensor): + return torch.add(tensor, tensor) + +def add_rref_to_value(rref, value): + return rref.to_here() + value + +class FaultyAgentRpcTest(RpcAgentTestFixture): + + # no faulty_messages defined so this fails all retryable messages - see + # faulty_rpc_agent_test_fixture.py for the list of retryable messages. + @dist_init(messages_to_delay={}) + def test_check_failed_messages(self): + if self.rank == 0: + dst_worker_b = worker_name((self.rank + 1) % self.world_size) + dst_worker_c = worker_name((self.rank + 2) % self.world_size) + + # Worker0 sends RPC to Worker1 and creates an RRef there + rref = rpc.remote(dst_worker_b, torch.add, args=(torch.ones(2, 2), torch.ones(2, 2))) + # Worker0 sends an RPC to Worker2 with the RRef as an arg + rpc.remote(dst_worker_c, add_rref_to_value, args=(rref, torch.ones(2, 2))) + # check if the output is as expected + self.assertEqual(rref.to_here(), torch.add(torch.ones(2, 2), torch.ones(2, 2))) + # explicitly delete all User RRefs + _delete_all_user_and_unforked_owner_rrefs() + + @dist_init + def test_verify_backend_options(self): + self.assertEqual(self.rpc_backend, rpc.backend_registry.BackendType.FAULTY_TENSORPIPE) + self.assertEqual(self.rpc_backend_options.num_worker_threads, 8) + self.assertEqual(self.rpc_backend_options.num_fail_sends, 3) + self.assertEqual(len(self.rpc_backend_options.messages_to_fail), 4) + self.assertEqual(len(self.rpc_backend_options.messages_to_delay), 2) + self.assertEqual(self.rpc_backend_options.rpc_timeout, rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + @dist_init(faulty_messages=["RREF_FORK_REQUEST", "RREF_CHILD_ACCEPT"]) + def test_custom_faulty_messages(self): + self.assertEqual( + {"RREF_FORK_REQUEST", "RREF_CHILD_ACCEPT"}, + set(self.rpc_backend_options.messages_to_fail), + ) + + @dist_init(faulty_messages=[]) + def test_no_faulty_messages(self): + self.assertEqual(len(self.rpc_backend_options.messages_to_fail), 0) + + @dist_init(messages_to_delay={"SCRIPT_CALL": 1.5}) + def test_custom_messages_to_delay(self): + self.assertEqual(self.rpc_backend_options.messages_to_delay, {"SCRIPT_CALL": 1.5}) + + def _test_remote_message_dropped_pickle(self, dst=None): + if self.rank != 0: + return + dst_rank = dst if dst is not None else (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + # Since we fail python_remote_call messages synchronously, the future + # corresponding to this remote call will be marked with an error when + # this function returns. + rref = rpc.remote(dst_worker, my_sleep_func, args=(1,)) + # Call to ensure pending callbacks are run. + wait_until_pending_futures_and_users_flushed() + # Attempt to fork the RRef should raise an error indicating the rpc.remote timeout. + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref._serialize() + # Test that using RRef as arg over RPC (which forks) results in the same + # error + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rpc.rpc_async(dst_worker, add_rref_to_value, args=(rref, 1)) + + @dist_init(faulty_messages=["PYTHON_REMOTE_CALL"]) + def test_remote_message_dropped_pickle(self): + self._test_remote_message_dropped_pickle() + + @dist_init(faulty_messages=["PYTHON_REMOTE_CALL"]) + def test_remote_message_dropped_pickle_to_self(self): + self._test_remote_message_dropped_pickle(self.rank) + + + def _test_remote_message_dropped_timeout(self, func, args, dst=None): + if self.rank != 0: + return + + # test the case where rpc.remote() message creation is completely dropped. + dst_rank = dst if dst is not None else (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + # Since we fail python_remote_call messages synchronously, the future + # corresponding to this remote call will be marked with an error when + # this function returns. + rref = rpc.remote(dst_worker, func, args=args) + # Call to ensure pending callbacks are run. + wait_until_pending_futures_and_users_flushed() + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref.to_here() + # Note: during shutdown, logs will indicate "Could not find OwnerRRef..." + # on the owning nodes, this is expected because the OwnerRRef was never + # successfully created. Therefore, delAllUsers will work as expected. + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_builtin_remote_message_dropped_timeout(self): + func = torch.add + args = (torch.tensor(1), torch.tensor(1)) + self._test_remote_message_dropped_timeout(func, args) + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_builtin_remote_message_dropped_timeout_to_self(self): + func = torch.add + args = (torch.tensor(1), torch.tensor(1)) + self._test_remote_message_dropped_timeout(func, args, dst=0) + + @dist_init(faulty_messages=["PYTHON_REMOTE_CALL"]) + def test_udf_remote_message_dropped_timeout(self): + func = my_sleep_func + args = (2,) + self._test_remote_message_dropped_timeout(func, args) + + @dist_init(faulty_messages=["PYTHON_REMOTE_CALL"]) + def test_udf_remote_message_dropped_timeout_to_self(self): + func = my_sleep_func + args = (2,) + self._test_remote_message_dropped_timeout(func, args, dst=0) + + def _test_remote_message_delay_timeout(self, func, args, dst=None): + if self.rank != 0: + return + # Test the case where remote message is eventually processed on the owner, + # but the future on the creator times out before the response comes back. + dst_rank = dst if dst is not None else (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + # 10 ms timeout + rref = rpc.remote(dst_worker, func, args=args, timeout=0.001) + # Future corresponding to the remote creation should time out. + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rref._get_future().wait() + + # Call to ensure pending callbacks are run. + wait_until_pending_futures_and_users_flushed() + # to_here() should now pick up that rpc.remote() creation has failed. + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref.to_here() + + # Test the case where rpc.remote() times out, but to_here() has already + # started blocking before. + # NOTE: we only test this when not sending to self, as to_here() calls + # calls localValue(), which does not send an RPC and thus does not have + # a timeout. This can be supported by allowing future.wait() to + # take in an optional timeout (https://github.com/pytorch/pytorch/issues/39280) + if dst_rank != self.rank: + slow_rref = rpc.remote(dst_worker, func, args=args, timeout=2) + + with self.assertRaisesRegex(RuntimeError, expected_error): + # to_here() should raise timeout error, since it does not know about the + # status of rpc.remote(). + slow_rref.to_here(0.001) + # Note: If we proceed with shutdown, UserRRef will send out a RRefUserDelete + # but this can be a noop since it may not exist on the owner yet. Later, + # the owner can process the RRef creation and wait for the delete message, + # thus leading to a timeout. + # Therefore, we wait until we get notification that pending owners have + # been confirmed before sending out RRefUserDeletes. + if dst_rank != self.rank: + wait_until_owners_and_forks_on_rank(2, 2, rank=dst_rank) + + @dist_init(faulty_messages=[], messages_to_delay={"PYTHON_REMOTE_CALL": 2}) + def test_udf_remote_message_delay_timeout(self): + func = my_sleep_func + args = (2,) + self._test_remote_message_delay_timeout(func, args) + + @dist_init(faulty_messages=[], messages_to_delay={"PYTHON_REMOTE_CALL": 2}) + def test_udf_remote_message_delay_timeout_to_self(self): + func = my_sleep_func + args = (1,) + self._test_remote_message_delay_timeout(func, args, dst=0) + + @dist_init( + faulty_messages=[], + messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1}, + ) + def test_remote_message_builtin_delay_timeout(self): + func = torch.add + args = (torch.tensor(1), torch.tensor(1)) + self._test_remote_message_delay_timeout(func, args) + + @dist_init( + faulty_messages=[], + messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1}, + ) + def test_remote_message_builtin_delay_timeout_to_self(self): + func = torch.add + args = (torch.tensor(1), torch.tensor(1)) + self._test_remote_message_delay_timeout(func, args, dst=0) + + @dist_init( + faulty_messages=[], + messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1}, + ) + def test_remote_message_script_delay_timeout(self): + func = my_script_func + args = (torch.tensor(1),) + self._test_remote_message_delay_timeout(func, args) + + @dist_init( + faulty_messages=[], + messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1}, + ) + def test_remote_message_script_delay_timeout_to_self(self): + func = my_script_func + args = (torch.tensor(1),) + self._test_remote_message_delay_timeout(func, args, dst=0) + + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_RREF_FETCH_CALL": 1}) + def test_rref_to_here_timeout(self): + if self.rank != 0: + return + + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rref.to_here(0.01) + + rref.to_here() + + @dist_init(faulty_messages=[]) + def test_rpc_builtin_timeout(self): + next_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(next_rank) + expected_error = self.get_timeout_error_regex() + # PYTHON_CALL message types which correspond to Python UDF over RPC + # by default get a delay (see faulty_rpc_agent_test_fixture) + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.rpc_sync( + dst_worker, + torch.add, + args=(torch.tensor(1), torch.tensor(1)), + timeout=1, + ) + + fut = rpc.rpc_async( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)), timeout=1 + ) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure that the currently set default timeout is large enough such + # that RPCs with delays still complete. + fut = rpc.rpc_async( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + fut.wait() + + # Ensure timeout if we set a new default and don't override + rpc._set_rpc_timeout(0.001) + fut = rpc.rpc_async( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure run to completion if we specify timeout of 0 + fut = rpc.rpc_async( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)), timeout=0 + ) + fut.wait() + # Reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_CALL": 1.5}) + def test_rpc_script_timeout(self): + next_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(next_rank) + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.rpc_sync(dst_worker, my_script_func, args=(torch.tensor(1),), timeout=1) + + fut = rpc.rpc_async(dst_worker, my_script_func, args=(torch.tensor(1),), timeout=1) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure that the currently set default timeout is large enough such + # that RPCs with delays still complete. + fut = rpc.rpc_async( + dst_worker, my_script_func, args=(torch.tensor(1),) + ) + fut.wait() + + # Ensure timeout if we set a new default and don't override + rpc._set_rpc_timeout(0.001) + fut = rpc.rpc_async( + dst_worker, my_script_func, args=(torch.tensor(1),) + ) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure run to completion if we specify timeout of 0 + rpc._set_rpc_timeout(0.001) + fut = rpc.rpc_async( + dst_worker, my_script_func, args=(torch.tensor(1),), timeout=0 + ) + fut.wait() + # Reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/faulty_rpc_agent_test_fixture.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/faulty_rpc_agent_test_fixture.py new file mode 100644 index 0000000000000000000000000000000000000000..ca584c1dc95aa0cd8f4d21ac8f0a31d214974106 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/faulty_rpc_agent_test_fixture.py @@ -0,0 +1,62 @@ +# mypy: allow-untyped-defs + +import torch.distributed.rpc as rpc +import torch.distributed.rpc._testing # noqa: F401 +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + +# The following message types are currently retried in the RREF protocol and +# distributed autograd. Thus only these messages should be tested with the +# Faulty RPC Agent. +retryable_message_types = ["RREF_FORK_REQUEST", + "RREF_CHILD_ACCEPT", + "RREF_USER_DELETE", + "CLEANUP_AUTOGRAD_CONTEXT_REQ"] + +# The following messages incur the corresponding delay in seconds while being +# processed in FaultyTensorPipeAgent's enqueueSend() function. +default_messages_to_delay = { + "PYTHON_CALL": 1.5, # Python UDF + "SCRIPT_CALL": 1.5, # Script/Builtin +} + +class FaultyRpcAgentTestFixture(RpcAgentTestFixture): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.messages_to_fail = retryable_message_types + self.messages_to_delay = default_messages_to_delay + + @property + def rpc_backend(self): + return rpc.backend_registry.BackendType[ + "FAULTY_TENSORPIPE" + ] + + @property + def rpc_backend_options(self): + return rpc.backend_registry.construct_rpc_backend_options( + self.rpc_backend, + init_method=self.init_method, + num_worker_threads=8, + num_fail_sends=3, + messages_to_fail=self.messages_to_fail, + messages_to_delay=self.messages_to_delay, + ) + + def setup_fault_injection(self, faulty_messages, messages_to_delay): + if faulty_messages is not None: + self.messages_to_fail = faulty_messages + if messages_to_delay is not None: + self.messages_to_delay = messages_to_delay + + def get_shutdown_error_regex(self): + error_regexes = [ + "Exception in thread pool task", + "Connection reset by peer", + "Connection closed by peer" + ] + return "|".join([f"({error_str})" for error_str in error_regexes]) + + def get_timeout_error_regex(self): + return "RPC ran for more than" diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__init__.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f41cccbc7307e22498bafc4419310480ed375d9c Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/dist_autograd_test.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/dist_autograd_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..130ca4907afeda7679629695df64540b9ee7c173 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/dist_autograd_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/rpc_test.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/rpc_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ced82b7fe60d9733c7ae15e9081e6faac685e1d9 Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/rpc_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/rpc_test_faulty.cpython-312.pyc b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/rpc_test_faulty.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f56099546c7f2270d9a3d9ec933135fc14b2cdee Binary files /dev/null and b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/__pycache__/rpc_test_faulty.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/dist_autograd_test.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/dist_autograd_test.py new file mode 100644 index 0000000000000000000000000000000000000000..0b69d9ff754481cf3ca771c561c9ecab22eb2cf8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/dist_autograd_test.py @@ -0,0 +1,115 @@ +# mypy: allow-untyped-defs + +from typing import Dict, Tuple + +import torch +import torch.distributed.autograd as dist_autograd +import torch.distributed.rpc as rpc +from torch import Tensor +from torch.distributed.rpc import rpc_async +from torch.testing import FileCheck +from torch.testing._internal.dist_utils import dist_init, worker_name +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +@torch.jit.script +def local_add(t1, t2): + return torch.add(t1, t2) + + +@torch.jit.script +def remote_add(t1, t2, dst: str): # noqa: E999 + return rpc_async(dst, local_add, (t1, t2)).wait() + + +@torch.jit.script +def fork_add(t1, t2, dst: str): + fut = torch.jit._fork(remote_add, t1, t2, dst) + return torch.jit._wait(fut) + + +class JitDistAutogradTest(RpcAgentTestFixture): + @dist_init + def test_get_gradients(self): + + @torch.jit.script + def dist_get_gradients(context_id: int) -> (Dict[Tensor, Tensor]): + return dist_autograd.get_gradients(context_id) + + FileCheck().check("get_gradients").run(str(dist_get_gradients.graph)) + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + t3 = torch.add(t1, t2) + + dist_autograd.backward(context_id, [t3.sum()]) + grads = dist_get_gradients(context_id) + + self.assertEqual(2, len(grads)) + self.assertIn(t1, grads) + self.assertIn(t2, grads) + self.assertEqual(torch.ones(3, 3), grads[t1]) + self.assertEqual(torch.ones(3, 3), grads[t2]) + + @dist_init + def test_dist_backward(self): + if self.rank != 0: + return + + @torch.jit.script + def dist_backward_script(context_id: int, loss: torch.Tensor): + dist_autograd.backward(context_id, [loss]) + + FileCheck().check("dist_backward").run(str(dist_backward_script.graph)) + with dist_autograd.context() as context_id: + t1 = torch.rand(3, 3, requires_grad=True) + t2 = torch.rand(3, 3, requires_grad=True) + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + loss = rpc.rpc_sync(dst_worker_name, torch.add, args=(t1, t2)).sum() + dist_backward_script(context_id, loss) + + @dist_init + def test_jit_fork_within_context(self): + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + res = fork_add(t1, t2, dst_worker_name) + loss = res.sum() + dist_autograd.backward(context_id, [loss]) + + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(2, len(grads)) + self.assertIn(t1, grads) + self.assertIn(t2, grads) + + @dist_init + def test_restore_context_after_swtich_to_jit_thread(self): + if self.rank != 0: + return + + @torch.jit.script + def forward_script( + context_id: int, dst_worker_name: str, t1: Tensor, t2: Tensor + ) -> Tuple[Tensor, Tensor]: + res1_fut = rpc.rpc_async(dst_worker_name, local_add, (t1, t1)) + res1 = res1_fut.wait() # After this, the script runs in a new JIT thread. + loss1 = res1.sum() + + # SendRpcBackward is not attached, since DistAutogradContext is lost here. + res2_fut = rpc.rpc_async(dst_worker_name, local_add, (t2, t2)) + res2 = res2_fut.wait() + loss2 = res2.sum() + + return loss1, loss2 + + with dist_autograd.context() as context_id: + t1 = torch.ones((2, 3), requires_grad=True) + t2 = torch.ones((2, 3), requires_grad=True) + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + loss0, loss1 = forward_script(context_id, dst_worker_name, t1, t2) + dist_autograd.backward(context_id, [loss0, loss1]) + grad0, grad1 = dist_autograd.get_gradients(context_id) + self.assertEqual(grad0, grad1) diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test.py new file mode 100644 index 0000000000000000000000000000000000000000..34c7fc793bf1945d622f08f0f2b106cf852034e8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test.py @@ -0,0 +1,1380 @@ +# mypy: allow-untyped-defs + +import time +import io +from typing import Dict, List, Tuple, Any + +import torch +import torch.distributed as dist +import torch.distributed.rpc as rpc +from torch import Tensor +from torch.autograd.profiler import record_function +from torch.distributed.rpc import RRef +from torch.distributed.rpc.internal import RPCExecMode, _build_rpc_profiling_key +from torch.futures import Future +from torch.testing._internal.common_utils import TemporaryFileName +from torch.testing._internal.dist_utils import ( + dist_init, + get_function_event, + initialize_pg, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + +from torch.autograd.profiler_legacy import profile as _profile + +def rref_isinstance(rref, cls_to_check): + return isinstance(rref.local_value(), cls_to_check) + +def sleep(t): + time.sleep(t) + + +def rpc_return_rref(dst): + return rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 1)) + + +@torch.jit.script +def rref_local_value(rref: RRef[Tensor]) -> Tensor: + return rref.local_value() + + +@torch.jit.script +def list_create() -> List[int]: + global_list = [1, 2, 3] + return global_list + + +@torch.jit.script +def rref_list_mutate(rref: RRef[List[int]]) -> None: + rref.local_value().append(4) + rref.to_here().append(5) + rref.to_here(5.0).append(6) + + +def return_value(value: int) -> int: + return value + + +class RRefAPITest: + @dist_init + def test_rref_is_owner(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + rref_var = rpc_return_rref(dst_worker_name) + + @torch.jit.script + def rref_tensor_is_owner(rref_var: RRef[Tensor]) -> bool: + return rref_var.is_owner() + + res = rref_tensor_is_owner(rref_var) + self.assertEqual(res, False) + + @dist_init + def test_rref_local_value(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + rref = rpc_return_rref(dst_worker_name) + + with self.assertRaisesRegex( + RuntimeError, r"Can't call RRef.local_value\(\) on a non-owner RRef" + ): + rref_local_value(rref) + + ret = ret = rpc.rpc_sync(dst_worker_name, rref_local_value, (rref,)) + self.assertEqual(ret, torch.add(torch.ones(2, 2), 1)) + + @dist_init + def test_local_rref_local_value(self): + if self.rank != 0: + return + + dst_worker_name = worker_name(self.rank) + rref = rpc.remote(dst_worker_name, return_value, (5,), {}) + + ret = rref_local_value(rref) + self.assertEqual(ret, 5) + + def _create_rref(self): + owner_rank = (self.rank + 2) % self.world_size + return rpc.remote( + worker_name(owner_rank), torch.add, args=(torch.zeros(2, 2), 1) + ) + + @dist_init + def test_user_rrefs_confirmed(self): + dst_rank = (self.rank + 1) % self.world_size + rref = self._create_rref() + ret = rpc.rpc_sync( + worker_name(dst_rank), script_check_rref_confirmed, args=(rref,) + ) + self.assertEqual(ret, True) + + @dist_init + def test_user_rrefs_confirmed_remote(self): + dst_rank = (self.rank + 1) % self.world_size + rref = self._create_rref() + ret_rref = rpc.remote( + worker_name(dst_rank), script_check_rref_confirmed, args=(rref,) + ) + self.assertEqual(ret_rref.to_here(), True) + + @dist_init + def test_rref_list_mutate(self): + dst = worker_name((self.rank + 1) % self.world_size) + list_rref = rpc.remote(dst, list_create) + + rpc.rpc_sync(dst, rref_list_mutate, args=(list_rref,)) + self.assertEqual(list_rref.to_here(), [1, 2, 3, 4, 5, 6]) + + +@torch.jit.script +def no_arg(): + return 0 + + +@torch.jit.script +def one_arg(value): + return value + 1 + +@torch.jit.script +def script_add_ones(x): + return torch.add(x, torch.ones(1)) + +@torch.jit.script +def script_add_ones_with_record_function(x, block: str): + with record_function(block): + return torch.add(x, torch.ones(1)) + + +@torch.jit.script +def record_function_on_caller_rpc_async(dst_worker_name: str, block: str) -> Tensor: + t: Tensor = torch.ones(1) + with record_function(block): + fut1 = rpc.rpc_async(dst_worker_name, script_add_ones, (t, )) + # Extra operator call to avoid de-duplication of the next async call + # see https://github.com/pytorch/pytorch/pull/62710#discussion_r694680279 + zero = torch.zeros_like(t) + fut2 = rpc.rpc_async(dst_worker_name, script_add_ones, (t, )) + res = fut1.wait() + fut2.wait() + zero + return res + + + +@torch.jit.script +def script_fork_wait_udf(tensor): + fut = torch.jit._fork(script_add_ones, tensor) + x = torch.jit._wait(fut) + return x + + +@torch.jit.script +def rref_to_here(rref_var: RRef[Tensor]) -> Tensor: + return rref_var.to_here() + + +@torch.jit.script +def return_rref(rref_var: RRef[Tensor]) -> RRef[Tensor]: + return rref_var + + +@torch.jit.script +def script_raise_func(value): + if value.numel() == 2: + raise ValueError("Expected error") + return value + 1 + + +@torch.jit.script +def script_fork_wait_throw(invalue): + fut = torch.jit._fork(script_raise_func, invalue) + value = torch.jit._wait(fut) + return value + + +@torch.jit.script +def call_rpc_with_profiling(record: torch.classes.profiler._RecordFunction, dst_worker_name: str) -> Tensor: + # Call rpc_async from within ScriptFunction and ensure that we can attach + # profiling callbacks. Note that handle here is a Tensor representation of + # RecordFunction. + fut = rpc.rpc_async(dst_worker_name, one_arg, (torch.tensor(1),)) + torch.ops.profiler._call_end_callbacks_on_jit_fut(record, fut) + ret = fut.wait() + return ret + +@torch.jit.script +def call_rpc_torchscript_with_record_function(dst_worker_name: str, block: str) -> Tensor: + fut = rpc.rpc_async(dst_worker_name, script_add_ones_with_record_function, (torch.tensor(1), block)) + return fut.wait() + + +@torch.jit.script +def call_fork_with_profiling(record: torch.classes.profiler._RecordFunction) -> Tensor: + # Call fork from within ScriptFunction and ensure that we can attach profiling + # callbacks to the resulting future. Note that handle here is a Tensor + # representation of RecordFunction. + fut = torch.jit._fork(one_arg, torch.tensor(1)) + torch.ops.profiler._call_end_callbacks_on_jit_fut(record, fut) + ret = fut.wait() + return ret + + +class MyScriptModuleWithRRefs(torch.jit.ScriptModule): + def __init__(self, dst_worker): + super().__init__() + self.rrefs = [] + for _ in range(4): + self.rrefs.append(rpc_return_rref(dst_worker)) + + @torch.jit.script_method + def forward(self) -> Tensor: + res_tensor = torch.ones(2, 2) + for rref in self.rrefs: + res_tensor += rref.to_here() + + return res_tensor + + +@torch.jit.ignore +def rref_python_annotation(rref_var: RRef[Tensor]) -> RRef[Tensor]: + return rref_var + + +@torch.jit.script +def rref_script_annotation(rref_var: RRef[Tensor]) -> Tensor: + return rref_python_annotation(rref_var).to_here() + + +class RRefTypingTest: + @dist_init + def test_rref_as_arg_and_return(self): + n = self.rank + 1 + dst_rank = n % self.world_size + local_ret = one_arg(torch.ones(2, 2)) + + # create rref on current rank + rref = rpc.remote(worker_name(self.rank), one_arg, args=(torch.ones(2, 2),)) + + # pass rref to another user in rpc call + ret = rpc.rpc_sync(worker_name(dst_rank), rref_to_here, args=(rref,)) + self.assertEqual(ret, local_ret) + + # return rref in rpc call + rref1 = rpc.rpc_sync(worker_name(dst_rank), return_rref, args=(rref,)) + self.assertEqual(rref1.to_here(), local_ret) + + # pass rref to another user in remote call + rref2 = rpc.remote(worker_name(dst_rank), rref_to_here, args=(rref,)) + self.assertEqual(rref2.to_here(), local_ret) + + # return rref in remote call + rref3 = rpc.remote(worker_name(dst_rank), return_rref, args=(rref,)) + self.assertEqual(rref3.to_here().to_here(), local_ret) + + @dist_init + def test_my_script_module_with_rrefs(self): + n = self.rank + 1 + dst_rank = n % self.world_size + + module_with_rrefs = MyScriptModuleWithRRefs(worker_name(dst_rank)) + res = module_with_rrefs() + self.assertEqual(res, torch.ones(2, 2) * 9) + + @dist_init + def test_rref_python_annotation(self): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_var = rpc_return_rref(worker_name(dst_rank)) + + res = rref_script_annotation(rref_var) + self.assertEqual(res, torch.ones(2, 2) + 1) + + +class FutureTypingTest: + @dist_init + def test_future_passed_between_python_and_jit(self): + dst_rank = (self.rank + 1) % self.world_size + inputs = (torch.tensor([1, 1]), torch.tensor([2, 2])) + ret_fut = rpc.rpc_async(worker_name(dst_rank), two_args_two_kwargs, args=inputs) + expected_res = torch.tensor([10, 10]) + + @torch.jit.script + def future_wait_in_script(fut: Future[Tensor]) -> Tensor: + return fut.wait() + + self.assertEqual(future_wait_in_script(ret_fut), expected_res) + + @torch.jit.script + def future_return_to_python( + dst_rank: int, inputs: Tuple[Tensor, Tensor] + ) -> Future[Tensor]: + return rpc.rpc_async( + f"worker{dst_rank}", two_args_two_kwargs, inputs + ) + + fut_res = future_return_to_python(dst_rank, inputs) + self.assertEqual(fut_res.wait(), expected_res) + + @dist_init + def test_future_python_annotation(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + input_0 = torch.ones(2, 2) + input_1 = 1 + expected_res = torch.add(input_0, input_1) + + @torch.jit.ignore + def python_return_future() -> Future[Tensor]: + fut = rpc.rpc_async(dst_worker_name, torch.add, (input_0, input_1), {}) + return fut + + @torch.jit.script + def script_use_future() -> Tensor: + fut = python_return_future() + return fut.wait() + + res = script_use_future() + self.assertEqual(res, expected_res) + + +@torch.jit.script +class MyScriptClass: + def __init__(self, a: int): + self.a = a + + def get_value(self) -> int: + return self.a + + +@torch.jit.interface +class MyModuleInterface(torch.nn.Module): + def forward(self) -> Tensor: + # pyre-ignore[7]: Pyre and torch.jit.interface don't mix well + pass + + +class MyScriptModule(torch.jit.ScriptModule): + def __init__(self, rank): + super().__init__() + self.a = torch.ones(rank) + + @torch.jit.script_method + def forward(self) -> Tensor: + return self.a + + @torch.jit.script_method + def custom_func(self) -> Tensor: + return self.a + + +def owner_create_rref_my_script_class(a): + return rpc.RRef(MyScriptClass(a)) + + +def owner_create_rref_my_script_module(a): + return rpc.RRef(MyScriptModule(a), type_hint=MyModuleInterface) + + +@torch.jit.script +def script_rref_get_value_my_script_class(rref: RRef[MyScriptClass]) -> int: + return rref.to_here().get_value() + + +@torch.jit.script +def script_rref_run_forward_my_script_module(rref: RRef[MyModuleInterface]) -> Tensor: + return rref.to_here().forward() + + +class LocalRRefTest: + @dist_init + def test_create_local_script_class_rref_in_py(self): + if self.rank != 0: + return + + # Create a local RRef. + rref_script_class = rpc.RRef(MyScriptClass(self.rank)) + ret = rref_script_class.to_here().get_value() + self.assertEqual(ret, self.rank) + + @dist_init + def test_create_local_script_module_rref_in_py(self): + if self.rank != 0: + return + + # Create a local RRef. + rref_script_module = rpc.RRef(MyScriptModule(self.rank), MyModuleInterface) + ret = rref_script_module.to_here().forward() + self.assertEqual(ret, torch.ones(self.rank)) + + # Create a local RRef without type hint. + with self.assertRaisesRegex( + RuntimeError, + ( + "The RRef being created contains a ScriptModule, " + "must provide its ModuleInterface type hint." + ), + ): + rref_script_module = rpc.RRef(MyScriptModule(self.rank)) + + @dist_init + def test_return_local_script_class_rref_in_py_and_use_in_script(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + # Create a local RRef remotely in Python. + rref = rpc.rpc_sync( + dst_worker_name, owner_create_rref_my_script_class, args=(self.rank,) + ) + + def use_rref_on_owner(rref: RRef[MyScriptClass]) -> int: + args = (rref,) + kwargs: Dict[str, Any] = {} + fut = rpc.rpc_async( + rref.owner(), script_rref_get_value_my_script_class, args, kwargs + ) + ret = fut.wait() + return ret + + # Use RRef in local Python RPC and remote Script run. + ret = use_rref_on_owner(rref) + self.assertEqual(ret, self.rank) + + # Use RRef in local Script RPC and remote Script run. + use_rref_on_owner_script = torch.jit.script(use_rref_on_owner) + ret = use_rref_on_owner_script(rref) + self.assertEqual(ret, self.rank) + + @dist_init + def test_return_local_script_module_rref_in_py_and_use_in_script(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + # Create a local RRef remotely in Python. + rref = rpc.rpc_sync( + dst_worker_name, owner_create_rref_my_script_module, args=(self.rank,) + ) + + def use_rref_on_owner(rref: RRef[MyModuleInterface]) -> Tensor: + args = (rref,) + kwargs: Dict[str, Any] = {} + fut = rpc.rpc_async( + rref.owner_name(), + script_rref_run_forward_my_script_module, + args, + kwargs, + ) + ret = fut.wait() + return ret + + # Use RRef in local Python RPC and remote Script run. + ret = use_rref_on_owner(rref) + self.assertEqual(ret, torch.ones(self.rank)) + + # Use RRef in local Script RPC and remote Script run. + use_rref_on_owner_script = torch.jit.script(use_rref_on_owner) + ret = use_rref_on_owner_script(rref) + self.assertEqual(ret, torch.ones(self.rank)) + + +def python_function(): + return 0 + + +@torch.jit.script +def two_args_two_kwargs( + first_arg, + second_arg, + first_kwarg=torch.tensor([3, 3]), + second_kwarg=torch.tensor([4, 4]), +): + return first_arg + second_arg + first_kwarg + second_kwarg + + +@torch.jit.script +def assorted_types_args_kwargs( + tensor_arg: Tensor, # noqa: E999 + str_arg: str, + int_arg: int, + tensor_kwarg: Tensor = torch.tensor([2, 2]), + str_kwarg: str = "str_kwarg", + int_kwarg: int = 2, +): + return tensor_arg + tensor_kwarg, str_arg + str_kwarg, int_arg + int_kwarg + + +@torch.jit.script +def raise_script(): + raise RuntimeError("Expected error") + + +@torch.jit.script +def script_rpc_async_call( + dst_worker_name: str, args: Tuple[Tensor, Tensor], kwargs: Dict[str, Tensor] +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + +@torch.jit.script +def script_rpc_sync_call( + dst_worker_name: str, args: Tuple[Tensor, Tensor], kwargs: Dict[str, Tensor] +): + res = rpc.rpc_sync(dst_worker_name, two_args_two_kwargs, args, kwargs) + return res + +@torch.jit.script +def script_rpc_remote_call( + dst_worker_name: str, args: Tuple[Tensor, Tensor], kwargs: Dict[str, Tensor] +): + rref_res = rpc.remote(dst_worker_name, two_args_two_kwargs, args, kwargs) + return rref_res.to_here() + +class JitRpcOpTest: + # Call functions remotely from Script. + @dist_init + def test_all_kwargs_are_populated_by_defaults(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = {} + + for script_op in [script_rpc_async_call, script_rpc_sync_call, script_rpc_remote_call]: + ret = script_op( + dst_worker_name, args, kwargs + ) + self.assertEqual(ret, torch.tensor([10, 10])) + + @dist_init + def test_some_kwargs_are_populated_by_defaults(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = {"first_kwarg": torch.tensor([2, 2])} + + for script_op in [script_rpc_async_call, script_rpc_sync_call, script_rpc_remote_call]: + ret = script_op( + dst_worker_name, args, kwargs + ) + self.assertEqual(ret, torch.tensor([9, 9])) + + @dist_init + def test_no_kwargs_are_populated_by_defaults(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = { + "first_kwarg": torch.tensor([2, 2]), + "second_kwarg": torch.tensor([3, 3]), + } + for script_op in [script_rpc_async_call, script_rpc_sync_call, script_rpc_remote_call]: + ret = script_op( + dst_worker_name, args, kwargs + ) + self.assertEqual(ret, torch.tensor([8, 8])) + + @dist_init + def test_args_and_kwargs_contain_different_types(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def script_rpc_async_call_with_assorted_types( + dst_worker_name: str, + ): + args = (torch.tensor([1, 1]), "str_arg", 1) + # Must annotate the value type as `Any`, because JIT type inference + # does not support multiple types when defining a Dict. + # The error JIT gives is, + # "Dict values must contain only a single type, " + # "expected: Tensor but found str instead." + kwargs: Dict[str, Any] = { + "tensor_kwarg": torch.tensor([3, 3]), + "str_kwarg": "_str_kwarg", + "int_kwarg": 3, + } + fut = rpc.rpc_async( + dst_worker_name, assorted_types_args_kwargs, args, kwargs + ) + ret = fut.wait() + return ret + + ret = script_rpc_async_call_with_assorted_types( + dst_worker_name + ) + self.assertEqual(ret, (torch.tensor([4, 4]), "str_arg_str_kwarg", 4)) + + @dist_init + def test_kwargs_not_passed(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def script_rpc_async_call_without_kwargs_passed( + dst_worker_name: str, + ): + args = () + fut = rpc.rpc_async(dst_worker_name, no_arg, args) + ret = fut.wait() + return ret + + ret = script_rpc_async_call_without_kwargs_passed( + dst_worker_name + ) + self.assertEqual(ret, 0) + + @dist_init + def test_args_kwargs_are_neither_passed(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def script_rpc_async_call_without_args_kwargs_passed( + dst_worker_name: str, + ): + fut = rpc.rpc_async(dst_worker_name, no_arg) + ret = fut.wait() + return ret + + ret = script_rpc_async_call_without_args_kwargs_passed( + dst_worker_name + ) + self.assertEqual(ret, 0) + + @dist_init + def test_less_than_needed_args_are_specified(self): + if self.rank != 0: + return + + # Notice, args matching happens during scripting. + with self.assertRaisesRegex(RuntimeError, "Argument second_arg not provided"): + + @torch.jit.script + def script_rpc_async_call_with_less_args( + dst_worker_name: str, # noqa: E999 + ): + args = (torch.tensor([1, 1]),) + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + + @dist_init + def test_more_than_needed_args_are_specified(self): + if self.rank != 0: + return + + # Notice, args matching happens during scripting. + with self.assertRaisesRegex( + RuntimeError, + "Expected at most 4 arguments but found 5 positional arguments", + ): + + @torch.jit.script + def script_rpc_async_call_with_more_args( + dst_worker_name: str, + ): + args = ( + torch.tensor([1, 1]), + torch.tensor([2, 2]), + torch.tensor([3, 3]), + torch.tensor([4, 4]), + torch.tensor([5, 5]), + ) + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + + @dist_init + def test_unexepected_kwarg_is_specified(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + # Notice, kwargs matching happens during execution. + @torch.jit.script + def script_rpc_async_call_with_unexpected_kwarg( + dst_worker_name: str, # noqa: E999 + ): + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = {"third_kwarg": torch.tensor([1, 1])} + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + + with self.assertRaisesRegex( + RuntimeError, "Unknown keyword argument 'third_kwarg'" + ): + ret = script_rpc_async_call_with_unexpected_kwarg( + dst_worker_name + ) + self.assertEqual(ret, 0) + + @dist_init + def test_call_python_function_remotely_from_script_not_supported(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def rpc_async_call_remote_py_function_in_torchscript(dst_worker_name: str): + args = () + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, python_function, args, kwargs) + ret = fut.wait() + return ret + + with self.assertRaisesRegex( + RuntimeError, "attempted to get undefined function" + ): + ret = rpc_async_call_remote_py_function_in_torchscript(dst_worker_name) + self.assertEqual(ret, 0) + + @dist_init + def test_call_script_function_that_raises_remotely_from_script(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + # Notice, TorchScript always translates(emits) Python `raise` statement, + # as the exception message string, "Exception", + # no matter what exception type and exception message are in the statement, + @torch.jit.script + def rpc_async_call_remote_raising_torchscript_in_torchscript( + dst_worker_name: str, + ): + args = () + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, raise_script, args, kwargs) + ret = fut.wait() + return ret + + with self.assertRaisesRegex(RuntimeError, "Expected error"): + ret = rpc_async_call_remote_raising_torchscript_in_torchscript( + dst_worker_name + ) + self.assertEqual(ret, 0) + + @dist_init + def test_call_script_function_that_not_exists_remotely_from_script(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def nonexisting_script(): + return 0 + + @torch.jit.script + def rpc_async_call_remote_nonexisting_torchscript_in_torchscript( + dst_worker_name: str, + ): + args = () + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, nonexisting_script, args, kwargs) + ret = fut.wait() + return ret + + with self.assertRaisesRegex( + RuntimeError, "attempted to get undefined function nonexisting_script" + ): + ret = rpc_async_call_remote_nonexisting_torchscript_in_torchscript( + dst_worker_name + ) + self.assertEqual(ret, 0) + + +@torch.jit.ignore +def my_script_module_init(rank: int) -> MyModuleInterface: + return MyScriptModule(rank) + + +@torch.jit.script +def construct_my_script_module(rank: int) -> MyModuleInterface: + return my_script_module_init(rank) + + +@torch.jit.script +def run_ref_script_module( + ref_script_module: RRef[MyModuleInterface], t: Tensor +) -> Tensor: + module = ref_script_module.to_here() + return module.forward() + t + + +@torch.jit.script +def script_check_rref_confirmed(rref: RRef[Tensor]) -> bool: + return rref.confirmed_by_owner() + + +@torch.jit.script +def save_rref(rref_var: RRef[Tensor], fname: str) -> None: + torch.save(rref_var, fname) + + +@torch.jit.script +def script_add(x: Tensor, y: Tensor) -> Tensor: + return x + y + + +@rpc.functions.async_execution +@torch.jit.script +def async_add(to: str, x: Tensor, y: Tensor) -> Future[Tensor]: + return rpc.rpc_async(to, script_add, (x, y)) + + +@rpc.functions.async_execution +@torch.jit.script +def async_wrong_type() -> Tensor: + return torch.zeros(2) + + +def load_script_module_with_pickled_rref(pickled_script_module): + f = io.BytesIO(pickled_script_module) + m = torch.jit.load(f) + return m() + + +class JitRpcTest( + RRefAPITest, + RRefTypingTest, + LocalRRefTest, + JitRpcOpTest, + FutureTypingTest, + RpcAgentTestFixture, +): + @dist_init + def test_torchscript_function(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + local_ret = one_arg(torch.ones(2, 2)) + ret = rpc.rpc_sync(dst_worker_name, one_arg, args=(torch.ones(2, 2),)) + self.assertEqual(ret, local_ret) + rref = rpc.remote(dst_worker_name, one_arg, args=(torch.ones(2, 2),)) + self.assertEqual(rref.to_here(), local_ret) + # create rref to itself + local_rref = rpc.remote( + worker_name(self.rank), one_arg, args=(torch.ones(2, 2),) + ) + self.assertEqual(local_rref.to_here(), local_ret) + + @dist_init + def test_torchscript_function_exception(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + with self.assertRaisesRegex(RuntimeError, r"one_arg\(\) expected at most"): + rpc.rpc_sync(dst_worker_name, one_arg, args=(10, 20)) + + with self.assertRaisesRegex(RuntimeError, r"one_arg\(\) expected at most"): + rpc.remote(dst_worker_name, one_arg, args=(10, 20)) + + @dist_init + def test_torchscript_functions_not_supported(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + my_local_script_module = MyScriptModule(self.rank) + + # It is not thread safe to instantiate MyScriptModule in multiple threads, + # wait for local MyScriptModule instantiation to finish, + # otherwise it could instantiate MyScriptModule in parallel with + # server thread in the below + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + + # rpc_sync still accepts script class and run it in + # the same code path as python call. + rpc.rpc_sync(dst_worker_name, MyScriptClass, args=(self.rank,)) + + # rpc_sync does not accept script module method. + # Python 3.5 and Python 3.6 throw different error message, the only + # common word can be greped is "pickle". + with self.assertRaisesRegex(TypeError, "pickle"): + rpc.rpc_async( + dst_worker_name, my_local_script_module.forward, args=() + ) + + @dist_init + def test_remote_script_module(self): + # TODO, need more investigation + # there is rref leak when shutting down, suspect it is because + # ref as arg is passed to pybind boundary, and the ref is not garbage + # collected by python when calling shutdown() + import torch.distributed.rpc.api as api + + api._ignore_rref_leak = True + + local_ret = torch.ones(self.rank) + torch.ones(self.rank) + + n = self.rank + 1 + dst_rank = n % self.world_size + remote_ref = rpc.remote( + worker_name(dst_rank), construct_my_script_module, args=(self.rank,) + ) + + # pass rref arg to owner + ret = rpc.rpc_sync( + worker_name(dst_rank), + run_ref_script_module, + args=(remote_ref, torch.ones(self.rank)), + ) + self.assertEqual(ret, local_ret) + + # pass rref arg to self/user + with self.assertRaisesRegex( + RuntimeError, + "is an RRef to a ScriptModule. It can't be sent through RPC from owner,", + ): + ret = rpc.rpc_sync( + worker_name(self.rank), + run_ref_script_module, + args=(remote_ref, torch.ones(self.rank)), + ) + + @dist_init + def test_create_script_module_on_remote(self): + dst_name = worker_name((self.rank + 1) % self.world_size) + # Construct on remote end with rpc_sync + created_script_module = rpc.rpc_sync( + dst_name, MyScriptModule, args=(self.rank,) + ) + # Forward should output a ones tensor of self.rank. + self.assertTrue(isinstance(created_script_module, torch.jit.ScriptModule)) + rank_ones_tensor = created_script_module() + self.assertEqual(torch.ones(self.rank), rank_ones_tensor) + + # Construct ScriptModule with rpc.remote. + remote_script_module = rpc.remote(dst_name, MyScriptModule, args=(self.rank,)) + # Verify it is an instance of ScriptModule on remote end. + remote_end_is_script = rpc.rpc_sync( + remote_script_module.owner(), + rref_isinstance, + args=(remote_script_module, torch.jit.ScriptModule), + ) + self.assertTrue(remote_end_is_script) + # Run forward pass remotely. + remote_forward_output = remote_script_module.rpc_sync().forward() + self.assertEqual(remote_forward_output, torch.ones(self.rank)) + # Run function defined on ScriptModule remotely. + remote_func_output = remote_script_module.rpc_sync().custom_func() + self.assertEqual(remote_func_output, torch.ones(self.rank)) + # Ensure we can transfer ScriptModule RRef to this rank and run + # forward pass. + local_script_module = remote_script_module.to_here() + self.assertTrue(isinstance(local_script_module, torch.jit.ScriptModule)) + rank_ones_tensor = local_script_module() + self.assertEqual(rank_ones_tensor, torch.ones(self.rank)) + local_script_func_output = local_script_module.custom_func() + self.assertEqual(local_script_func_output, torch.ones(self.rank)) + + @dist_init + def test_load_script_module_with_pickled_rref(self): + dst_name = worker_name((self.rank + 1) % self.world_size) + m1 = MyScriptModuleWithRRefs(dst_name) + m2 = MyScriptModuleWithRRefs(dst_name) + + f = io.BytesIO() + + rpc._enable_jit_rref_pickle() + torch.jit.save(m1, f) + rpc._disable_jit_rref_pickle() + + out1 = rpc.rpc_sync( + dst_name, + load_script_module_with_pickled_rref, + args=(f.getvalue(),) + ) + out2 = m2() + self.assertEqual(out1, out2) + + @dist_init + def test_rref_jit_pickle_not_supported(self): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_var = rpc_return_rref(worker_name(dst_rank)) + with TemporaryFileName() as fname: + with self.assertRaisesRegex( + RuntimeError, "RRef jit pickling is only allowed inside RPC calls" + ): + save_rref(rref_var, fname) + + @dist_init + def test_remote_script_throw(self): + rref = rpc.remote( + worker_name((self.rank + 1) % self.world_size), + script_raise_func, + args=(torch.ones(2),), + ) + with self.assertRaisesRegex(Exception, ".*Expected error.*"): + rref.to_here() + + @dist_init + def test_remote_script_udf(self): + rref = rpc.remote( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_udf, + args=(torch.ones(2),), + ) + self.assertEqual(rref.to_here(), torch.ones(2) * 2) + + @dist_init + def test_async_script_udf(self): + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_udf, + args=(torch.ones(2),), + ) + self.assertEqual(future.wait(), torch.ones(2) * 2) + + @dist_init + def test_callback_simple(self): + def callback(fut): + return fut.wait() + 1 + + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_udf, + args=(torch.ones(2),), + ).then(callback) + self.assertEqual(future.wait(), torch.ones(2) * 2 + 1) + + @dist_init + def test_callback_chain(self): + n = self.rank + 1 + + def callback(fut): + return fut.wait() + 1 + + fut = rpc.rpc_async( + worker_name(n % self.world_size), one_arg, args=(torch.ones(n, n),) + ) + + num_cbs = 20 + for _ in range(num_cbs): + fut = fut.then(callback) + + self.assertEqual(fut.wait(), torch.ones(n, n) + 1 + num_cbs) + + @dist_init + def test_add_done_callback(self): + callback_called = None + + def callback(fut): + nonlocal callback_called + callback_called = fut.wait() * 2 + + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_udf, + args=(torch.ones(2),), + ) + + future.add_done_callback(callback) + future_then = future.then(lambda _: True) + + self.assertEqual(future.wait(), torch.ones(2) * 2) + + # We have no guarantee that the add_done_callback fn will execute before the test finishes. + # Adding a 'then' callback that runs afterwards to guarantee we wait for the first callback + future_then.wait() + self.assertEqual(callback_called, torch.ones(2) * 4) + + @dist_init + def test_async_script_throw(self): + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_throw, + args=(torch.ones(2),), + ) + with self.assertRaisesRegex(Exception, ".*Expected error.*"): + future.wait() + + @dist_init + def test_callback_with_exception(self): + def callback(fut): + with self.assertRaisesRegex(Exception, ".*Expected error.*"): + fut.wait() + raise RuntimeError("Another expected error") + + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_throw, + args=(torch.ones(2),), + ).then(callback) + + with self.assertRaisesRegex(RuntimeError, "Another expected error"): + future.wait() + + @dist_init + def test_call_rpc_with_profiling(self): + # Ensures that we can call torch.ops.profiler._call_end_callbacks_on_jit_fut on a jit + # future from within a script function that calls rpc_async + if self.rank == 0: + with _profile() as prof: + prof_key = _build_rpc_profiling_key( + RPCExecMode.ASYNC, + torch._jit_internal._qualified_name(one_arg), + "worker0", + "worker1", + ) + with torch.autograd.profiler.record_function(prof_key) as rf: + call_rpc_with_profiling(rf.record, "worker1") + # TODO: Can't get a reliable time for this profiling event since + # it's hard to estimate the execution time on the remote end for non-UDFs. + # This can be resolved by https://github.com/pytorch/pytorch/issues/36272. + # After that, this test should be modified to validate the function time. + events = prof.function_events + function_event = get_function_event(events, prof_key) + self.assertTrue(torch._jit_internal._qualified_name(one_arg) in function_event.name) + + @dist_init + def test_rpc_async_jit_profiled(self): + # Tests that rpc_async calls made from within a TorchScript function are + # profiled. + if self.rank == 0: + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = worker_name(dst_rank) + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = {} + with _profile() as prof: + script_rpc_async_call( + dst_worker_name, args, kwargs + ) + + # Ensure rpc_async call is profiled + function_events = prof.function_events + qual_name = torch._jit_internal._qualified_name(two_args_two_kwargs) + rpc_async_jit_event = [ + event + for event in function_events + if qual_name in event.name and event.node_id == self.rank + ] + self.assertEqual(len(rpc_async_jit_event), 1) + rpc_async_jit_event = rpc_async_jit_event[0] + profiled_name = _build_rpc_profiling_key( + RPCExecMode.ASYNC_JIT, + qual_name, + worker_name(self.rank), + dst_worker_name, + ) + self.assertEqual(profiled_name, rpc_async_jit_event.name) + remote_events = [event for event in function_events if event.is_remote] + # All remote events should have taken place on dst_rank + remote_event_node_ids = { + remote_event.node_id for remote_event in remote_events + } + self.assertEqual(remote_event_node_ids, {dst_rank}) + # script_rpc_async_call invokes add operator + # so we should see this as a remote event. + remote_add = next( + remote_event + for remote_event in remote_events + if "aten::add" in remote_event.name + ) + remote_add_profiled_name = f"{profiled_name}#remote_op: aten::add" + self.assertEqual(remote_add.name, remote_add_profiled_name) + + @dist_init + def test_record_function_on_caller_rpc_async(self): + if self.rank == 0: + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = worker_name(dst_rank) + block_scope = "foo" + with _profile() as prof: + # Runs 2 rpc_async calls within JIT under record_function. + record_function_on_caller_rpc_async(dst_worker_name, block_scope) + + # Ensure record_function event is profiled. + function_events = prof.function_events + record_function_scope_event = [ + event for event in function_events if event.name == block_scope + ] + self.assertEqual(1, len(record_function_scope_event)) + record_function_scope_event = record_function_scope_event[0] + # Ensure RPC future is profiled. + expected_key = _build_rpc_profiling_key( + RPCExecMode.ASYNC_JIT, + torch._jit_internal._qualified_name(script_add_ones), + worker_name(self.rank), + dst_worker_name, + ) + jit_rpc_events = [ + event for event in function_events if event.name == expected_key + ] + self.assertEqual(2, len(jit_rpc_events)) + # Validate that the record_function scope time is greater than both + # of the individual RPC async call times. The reason it is not necessarily + # greater than the sum is because the two can execute in parallel. + for jit_rpc_event in jit_rpc_events: + self.assertTrue( + record_function_scope_event.cpu_time_total + > jit_rpc_event.cpu_time_total + ) + + @dist_init + def test_rpc_torchscript_record_function(self): + # tests that torchscript functions can be profiled using with + # record_function(...) over RPC. + REMOTE_OP_STR = "#remote_op: " + if self.rank == 0: + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = worker_name(dst_rank) + block_scope = "foo" + with _profile() as prof: + call_rpc_torchscript_with_record_function(dst_worker_name, block_scope) + + # Need to call below to populate CPU children. + prof.key_averages() + function_events = prof.function_events + expected_key = ( + _build_rpc_profiling_key( + RPCExecMode.ASYNC_JIT, + torch._jit_internal._qualified_name( + script_add_ones_with_record_function + ), + worker_name(self.rank), + dst_worker_name, + ) + + REMOTE_OP_STR + + block_scope + ) + remote_record_function_event = next( + evt for evt in function_events if evt.name == expected_key + ) + self.assertTrue(block_scope in remote_record_function_event.name) + remote_children = remote_record_function_event.cpu_children + self.assertTrue("aten::add" in child.name for child in remote_children) + + def test_record_function_jit_end_callbacks_with_fork(self): + # Ensures that we can call rf._call_end_callbacks_on_future on a jit + # future in python eager mode with torch.jit.fork + sleep_interval = 1 + with _profile() as prof: + with torch.autograd.profiler.record_function("foo") as rf: + fut = torch.jit._fork(sleep, sleep_interval) + rf._call_end_callbacks_on_future(fut) + fut.wait() + + function_events = prof.function_events + sleep_event = get_function_event(function_events, "foo") + self.assertEqual(sleep_event.name, "foo") + # Validate that callbacks were fired at the right time by checking the + # profiling event cpu time + self.assertGreaterAlmostEqual(sleep_event.cpu_time * 1e-6, sleep_interval) + + def test_call_fork_in_jit_with_profiling(self): + # Ensures that we can call torch.ops.profiler._call_end_callbacks_on_jit_fut on a jit + # future from within a script function with torch.jit.fork + with _profile() as prof: + with torch.autograd.profiler.record_function("foo") as rf: + call_fork_with_profiling(rf.record) + + events = prof.function_events + function_event = get_function_event(events, "foo") + self.assertEqual(function_event.name, "foo") + + @dist_init + def test_async_function_simple(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + ret = rpc.rpc_sync( + dst1, async_add, args=(dst2, torch.ones(2, 2), torch.ones(2, 2)) + ) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + @dist_init + def test_async_function_wrong_return_type(self): + with self.assertRaisesRegex( + RuntimeError, + "Async functions must return an IValue of Future type, but got Tensor", + ): + rpc.rpc_sync( + worker_name((self.rank + 1) % self.world_size), async_wrong_type + ) + + @dist_init + def test_async_function_wrong_decorator_order(self): + # @torch.jit.script complains about undefined value rpc. Error is shown + # below. The reason for not checking error string is to avoid making + # JIT error handling code depend on RPC tests, as we don't have any + # restrictions on the error message here. + # + # RuntimeError: + # undefined value rpc: + # def async_wrong_decorator_order(to, x, y): + # # type: (str, Tensor, Tensor) -> Future[Tensor] + # return rpc.rpc_async(to, script_add, (x, y)) + # ~~~ <--- HERE + with self.assertRaises(RuntimeError): + + @torch.jit.script + @rpc.functions.async_execution + def async_wrong_decorator_order( + to: str, x: Tensor, y: Tensor + ) -> Future[Tensor]: + return rpc.rpc_async(to, script_add, (x, y)) + + @dist_init + def test_async_function_remote(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + rref = rpc.remote( + dst1, async_add, args=(dst2, torch.ones(2, 2), torch.ones(2, 2)) + ) + self.assertEqual(rref.to_here(), torch.ones(2, 2) + 1) + + @dist_init + def test_async_function_remote_multi(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + num = 20 + rrefs = [ + rpc.remote( + dst1, + async_add, + args=(dst2, torch.ones(2, 2), torch.ones(2, 2) * i) + ) for i in range(num) + ] + + for i in range(num): + self.assertEqual(rrefs[i].to_here(), torch.ones(2, 2) + i) + + @dist_init + def test_async_function_wrong_return_type_remote(self): + rref = rpc.remote( + worker_name((self.rank + 1) % self.world_size), async_wrong_type + ) + + with self.assertRaisesRegex( + RuntimeError, + "Async functions must return an IValue of Future type, but got Tensor", + ): + rref.to_here() diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test_faulty.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test_faulty.py new file mode 100644 index 0000000000000000000000000000000000000000..0e49793634197e845969c745083d2532ab2988f4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test_faulty.py @@ -0,0 +1,218 @@ +# mypy: allow-untyped-defs + +from typing import Dict, Tuple + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch.distributed.rpc import RRef +from torch.testing._internal.dist_utils import ( + dist_init, + worker_name, + wait_until_pending_futures_and_users_flushed +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +@torch.jit.script +def two_args_two_kwargs( + first_arg, + second_arg, + first_kwarg=torch.tensor([3, 3]), + second_kwarg=torch.tensor([4, 4]), +): + return first_arg + second_arg + first_kwarg + second_kwarg + + +@torch.jit.script +def script_rpc_async_call( + dst_worker_name: str, args: Tuple[Tensor, Tensor], kwargs: Dict[str, Tensor] +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + + +@torch.jit.script +def rpc_async_call_with_timeout( + dst_worker_name: str, + args: Tuple[Tensor, Tensor], + kwargs: Dict[str, Tensor], + timeout: float, +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs, timeout) + ret = fut.wait() + return ret + + +@torch.jit.script +def rpc_async_call_with_timeout_future_ret( + dst_worker_name: str, + args: Tuple[Tensor, Tensor], + kwargs: Dict[str, Tensor], + timeout: float, +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs, timeout) + return fut + + +@torch.jit.script +def rpc_async_call_future_ret( + dst_worker_name: str, args: Tuple[Tensor, Tensor], kwargs: Dict[str, Tensor] +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + return fut + +@torch.jit.script +def rref_to_here(rref_var: RRef[Tensor]) -> Tensor: + return rref_var.to_here() + +@torch.jit.script +def rref_to_here_with_timeout(rref_var: RRef[Tensor], timeout: float) -> Tensor: + return rref_var.to_here(timeout) + +@torch.jit.script +def rpc_async_with_rref_arg(dst_worker_name: str, args: Tuple[RRef[Tensor]]) -> Tensor: + fut = rpc.rpc_async(dst_worker_name, rref_to_here, args) + ret = fut.wait() + return ret + + +class JitFaultyAgentRpcTest(RpcAgentTestFixture): + """ + Run tests for rpc_async in JIT under the faulty agent test fixture to test + arbitrary timeouts. + """ + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_CALL": 1.5}) + def test_timeout_in_torchscript_function(self): + # Call rpc_async + fut.wait() in torchscript function and ensure that + # timeout is raised. + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = { + "first_kwarg": torch.tensor([2, 2]), + "second_kwarg": torch.tensor([3, 3]), + } + expected_error = self.get_timeout_error_regex() + # Ensure that we get a timeout if we override the default timeout and + # the RPC takes longer to execute. + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc_async_call_with_timeout(dst_worker_name, args, kwargs, 0.5) + + # Ensure that we timeout if we don't specify a timeout but the default + # is less than the RPC takes to execute. + rpc._set_rpc_timeout(0.001) + with self.assertRaisesRegex(RuntimeError, expected_error): + script_rpc_async_call( + dst_worker_name, args, kwargs + ) + + # Ensure that we run to completion if zero timeout is specified. + ret = rpc_async_call_with_timeout(dst_worker_name, args, kwargs, 0) + self.assertEqual(ret, torch.tensor([8, 8])) + # reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_CALL": 1.5}) + def test_timeout_in_python(self): + # Ensures timeouts are raised if we call rpc_async from within a + # torchscript function, but wait on the future in python. + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = { + "first_kwarg": torch.tensor([2, 2]), + "second_kwarg": torch.tensor([3, 3]), + } + expected_error = self.get_timeout_error_regex() + + fut = rpc_async_call_with_timeout_future_ret(dst_worker_name, args, kwargs, 0.5) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure timeout if we don't specify but the default is less than the + # RPC takes to execute. + rpc._set_rpc_timeout(0.001) + fut = rpc_async_call_future_ret(dst_worker_name, args, kwargs) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure run to completion if zero timeout is specified + fut = rpc_async_call_with_timeout_future_ret(dst_worker_name, args, kwargs, 0) + result = fut.wait() + self.assertEqual(result, torch.tensor([8, 8])) + # reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_remote_timeout_to_here_in_jit(self): + # Test that calling to_here() in JIT will raise timeout error if + # rpc.remote failed. + if self.rank != 0: + return + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + # Will ensure error handling callbacks are run. + wait_until_pending_futures_and_users_flushed() + # Call to_here() within a ScriptFunction and ensure it raises + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref_to_here(rref) + + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_RREF_FETCH_CALL": 1}) + def test_rref_to_here_timeout_in_jit(self): + if self.rank != 0: + return + + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rref_to_here_with_timeout(rref, 0.01) + + rref_to_here_with_timeout(rref, 100) + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_rref_timeout_pickle_in_jit(self): + if self.rank != 0: + return + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + # Will ensure error handling callbacks are run. + wait_until_pending_futures_and_users_flushed() + # Call RPC with RRef arg in JIT, which will go through JIT pickling and + # ensure error is raised. + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rpc_async_with_rref_arg(dst_worker, (rref, )) + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_rref_timeout_pickle_script_func(self): + # Similar to above test, but calls python rpc with script function. + if self.rank != 0: + return + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + # Will ensure error handling callbacks are run. + wait_until_pending_futures_and_users_flushed() + # Call RPC with script function that takes RRef, ensure timeout during pickling + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rpc.rpc_sync(dst_worker, rref_to_here, args=(rref, )) diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/rpc_agent_test_fixture.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/rpc_agent_test_fixture.py new file mode 100644 index 0000000000000000000000000000000000000000..3a684b73d2f315a00465371fad3050a795251ddb --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/rpc_agent_test_fixture.py @@ -0,0 +1,63 @@ +# mypy: allow-untyped-defs + +import os +from abc import ABC, abstractmethod + +import torch.testing._internal.dist_utils + + +class RpcAgentTestFixture(ABC): + @property + def world_size(self) -> int: + return 4 + + @property + def init_method(self): + use_tcp_init = os.environ.get("RPC_INIT_WITH_TCP", None) + if use_tcp_init == "1": + master_addr = os.environ["MASTER_ADDR"] + master_port = os.environ["MASTER_PORT"] + return f"tcp://{master_addr}:{master_port}" + else: + return self.file_init_method + + @property + def file_init_method(self): + return torch.testing._internal.dist_utils.INIT_METHOD_TEMPLATE.format( + file_name=self.file_name + ) + + @property + @abstractmethod + def rpc_backend(self): + pass + + @property + @abstractmethod + def rpc_backend_options(self): + pass + + def setup_fault_injection(self, faulty_messages, messages_to_delay): # noqa: B027 + """Method used by dist_init to prepare the faulty agent. + + Does nothing for other agents. + """ + + # Shutdown sequence is not well defined, so we may see any of the following + # errors when running tests that simulate errors via a shutdown on the + # remote end. + @abstractmethod + def get_shutdown_error_regex(self): + """ + Return various error message we may see from RPC agents while running + tests that check for failures. This function is used to match against + possible errors to ensure failures were raised properly. + """ + + @abstractmethod + def get_timeout_error_regex(self): + """ + Returns a partial string indicating the error we should receive when an + RPC has timed out. Useful for use with assertRaisesRegex() to ensure we + have the right errors during timeout. + """ diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/rpc_test.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/rpc_test.py new file mode 100644 index 0000000000000000000000000000000000000000..a461c3a86bf13c50e4a258434a641aa9665dc8a2 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/rpc_test.py @@ -0,0 +1,6473 @@ +# mypy: allow-untyped-defs + +import concurrent.futures +import contextlib +import json +import os +import sys +import threading +import time + +from collections import namedtuple +from functools import partial +from threading import Event +from threading import Lock +from unittest import mock + +import torch +import torch.nn as nn +import torch.distributed as dist +import torch.distributed.rpc as rpc +import torch.distributed.autograd as dist_autograd +from torch.distributed.rpc import RRef, _get_debug_info, _rref_context_get_debug_info, WorkerInfo +from torch.distributed.rpc.api import _use_rpc_pickler, _thread_local_var, _wait_all +from torch.distributed.rpc.internal import ( + PythonUDF, + RPCExecMode, + _internal_rpc_pickler, + _build_rpc_profiling_key, +) +from torch.futures import Future +from torch.testing._internal.common_distributed import ( + skip_if_lt_x_gpu, + captured_output, + tp_transports, +) +from torch.testing._internal.common_utils import ( + IS_MACOS, + load_tests, + skip_but_pass_in_sandcastle_if, + get_cycles_per_ms, +) + +from torch.testing._internal.dist_utils import ( + dist_init, + get_function_event, + initialize_pg, + wait_until_node_failure, + wait_until_pending_futures_and_users_flushed, + wait_until_owners_and_forks_on_rank, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) +from torch.testing._internal.common_utils import TemporaryFileName + +from torch.autograd.profiler_legacy import profile as _profile +import operator + + +def foo_add(): + return torch.add(torch.ones(1), torch.ones(1)) + +def udf_with_torch_ops(device=-1, use_record_function=False): + device_ctx = contextlib.nullcontext() if device == -1 else torch.cuda.device(device) + record_function_ctx = ( + torch.autograd.profiler.record_function("##forward##") + if use_record_function + else contextlib.nullcontext() + ) + with device_ctx, record_function_ctx: + t1, t2 = torch.ones(1), torch.ones(1) + t = torch.add(t1, t2) + t = torch.mul(t, t) + t = t.relu() + t = t.sigmoid() + +# Events (operator invocations) that are expected to be ran as part of the above +# function. +EXPECTED_REMOTE_EVENTS = [ + "aten::ones", + "aten::ones", + "aten::add", + "aten::mul", + "aten::relu", + "aten::clamp_min", + "aten::sigmoid", +] + +# Remote operations are prefixed with the following string for RPC profiling. +REMOTE_OP_STR = "#remote_op: " + + +VALUE_FUTURE = concurrent.futures.Future() +DONE_FUTURE = concurrent.futures.Future() + +FIFTY_MIL_CYCLES = 50000000 + +_rpc_barrier_count = 0 + +def _increment_count(): + global _rpc_barrier_count + _rpc_barrier_count += 1 + +def _reset_count(): + global _rpc_barrier_count + _rpc_barrier_count = 0 + +class StubRpcAgent: + def __init__(self, world_size): + self.world_size = world_size + + def get_worker_infos(self): + return { + WorkerInfo(name=worker_name(rank), id=rank) + for rank in range(self.world_size) + } + + +def _stub_construct_rpc_backend_options_handler(**kwargs): + return mock.Mock() # RpcBackendOptions. + + +def _stub_init_rpc_backend_handler(store, name, rank, world_size, rpc_backend_options): + return StubRpcAgent(world_size=world_size) + + +def set_value(value): + VALUE_FUTURE.set_result(value) + + +def wait_for_value_future(): + return VALUE_FUTURE.result() + + +def set_and_check_done(value): + VALUE_FUTURE.set_result(value) + return DONE_FUTURE.result() + + +# it is used to test python user defined function over rpc +# classes and functions are used to test python user defined class and +# methods over rpc +TensorClass = namedtuple("TensorClass", ["tensors"]) + +class MyPickleClass: + def __init__(self) -> None: + self.t = None + + def __getstate__(self): + (pickled_python_udf, tensors) = _internal_rpc_pickler.serialize( + PythonUDF(my_tensor_function, (torch.ones(2, 2), torch.ones(2, 2)), None) + ) + return (pickled_python_udf, tensors) + + def __setstate__(self, obj): + python_udf = _internal_rpc_pickler.deserialize(obj[0], obj[1]) + result = python_udf.func(python_udf.args[0], python_udf.args[1]) + self.t = result + + def set(self, val): + self.t = val + + +class SlowPickleClass: + def __init__(self, t): + self.t = t + + def __getstate__(self): + time.sleep(self.t) + return (self.t, ) + + def __setstate__(self, obj): + self.t = obj[0] + time.sleep(self.t) + + +class MyClass: + def __init__(self, a, delay=False): + self.a = a + # delay initialization to simulate errors if specified + if delay: + time.sleep(2) + + def my_instance_method(self, b): + return self.a + b + + @classmethod + def my_class_method(cls, d, e): + return d + e + + @staticmethod + def my_static_method(f): + return f > 10 + + def increment_value(self, increment): + self.a += increment + + def get_value(self): + return self.a + + def my_slow_method(self, my_tensor_arg): + time.sleep(5) + return torch.add(self.a, my_tensor_arg) + + +def _call_method_on_rref(method, rref, *args, **kwargs): + return method(rref.local_value(), *args, **kwargs) + + +def get_rref_list(values): + return [RRef(MyClass(a)) for a in values] + + +def add_rref_to_value(rref, value): + return rref.to_here() + value + + +def run_nested_pickle(pickle_cls_instance, tensor): + return pickle_cls_instance.t + tensor + +def build_sparse_tensor(coalesce=False): + i = [[0, 1, 1], [2, 0, 2]] + v = [3, 4, 5] + tensor = torch.sparse_coo_tensor(i, v, (2, 3)) + if coalesce: + tensor = tensor.coalesce() + return tensor + +def build_complex_tensors(): + a = torch.ones(3, 3) + b = [a, a] + c = [b, b] + d = [a, b] + e = {a: d} + return [a, b, c, d, e] + +def non_cont_test(t_view, t_cont): + if t_view.is_contiguous(): + raise Exception('t_view is contiguous!') # noqa: TRY002 + if not t_cont.is_contiguous(): + raise Exception('t_cont is not contiguous!') # noqa: TRY002 + if not torch.equal(t_view, t_cont): + raise Exception('t_view is not equal to t_cont!') # noqa: TRY002 + return t_view + +def my_function(a, b, c): + return a + b + c + + +def my_tensor_function(a, b): + return a + b + +def my_container_sum(a): + result = a[0] + for tensor in a[1:]: + result += tensor + return result + + +def my_sleep_func(seconds=1): + time.sleep(seconds) + return torch.mul(torch.tensor(1), torch.tensor(1)) + + +def my_complex_tensor_function(list_input, tensor_class_input, dict_input): + res = list_input[0] + for t in list_input: + res += t + for v in dict_input.values(): + res += v + complex_tensors = tensor_class_input.tensors + return (res, complex_tensors[0], complex_tensors[1], complex_tensors[2]) + + +def my_rref_function(rref_a, rref_b): + return rref_a.to_here() + rref_b.to_here() + + +def delayed_add(a, b, seconds=0.05): + time.sleep(seconds) + return a + b + + +def identity(a): + return a + +def no_result(): + print("do nothing") + +def raise_or_inc(value): + if value.numel() == 2: + raise ValueError("Expected error") + return value + 1 + +def nested_rpc(dst): + return rpc.rpc_sync(dst, torch.add, args=(torch.ones(2, 2), 1)) + + +def nested_rpc_sparse(dst): + return rpc.rpc_sync( + dst, + torch.add, + args=(build_sparse_tensor(), build_sparse_tensor()) + ) + + +def multi_layer_nested_async_rpc(dst, world_size, ttl): + # this method returns immediately without blocking the callee, but will + # generate additional requests. + if ttl > 0: + current_dst = worker_name(dst) + next_dst = (dst + 1) % world_size + rpc.rpc_async( + current_dst, + multi_layer_nested_async_rpc, + args=(next_dst, world_size, ttl - 1), + ) + return 0 + + +def nested_rref(dst): + return ( + rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 1)), + rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 2)), + ) + + +def nested_rref_sparse(dst): + return ( + rpc.remote( + dst, + torch.add, + args=(build_sparse_tensor(), build_sparse_tensor()) + ), + rpc.remote( + dst, + torch.add, + args=(build_sparse_tensor(), build_sparse_tensor()) + ), + ) + + +def nested_remote(dst): + rref = rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 3)) + return rref.to_here() + +def nested_remote_sparse(dst): + rref = rpc.remote(dst, torch.add, args=(build_sparse_tensor(), build_sparse_tensor())) + return rref.to_here() + + +def rref_forward_chain(dst, world_size, rref, ttl): + if ttl > 0: + current_dst = worker_name(dst) + next_dst = (dst + 1) % world_size + ret_rref = rpc.remote( + current_dst, rref_forward_chain, args=(next_dst, world_size, rref, ttl - 1) + ) + return [ret_rref] + else: + return rref.to_here() + + +def rpc_return_rref(dst): + return rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 1)) + + +def light_rpc(): + return 0 + + +def heavy_rpc(tensor): + for i in range(1, 100): + tensor *= i + tensor /= i + 1 + return 0 + + +def heavy_rpc_sparse(tensor): + for i in range(1, 100): + tensor *= i + tensor = tensor / (i + 1) + return 0 + +@torch.jit.script +def heavy_rpc_torchscript(tensor): + for i in range(1, 100): + tensor *= i + tensor /= i + 1 + return 0 + + +@torch.jit.script +def my_script_func(tensor): + return torch.add(tensor, tensor) + + +expected_err = "Expected error" + +# Note that it needs to inherit from Exception, not BaseException. See comment +# in rpc/internal.py +class CustomException(Exception): + def __init__(self, bool, msg): + self.bool = bool + super().__init__(msg) + +def raise_func(): + raise ValueError(expected_err) + +def custom_raise_func(): + raise CustomException(True, "foo") + +@torch.jit.script +def raise_func_script(expected_err: str) -> torch.Tensor: + raise ValueError(expected_err) + +expected_err_escape = "\nFirst line of error \n next line of error \n last line of error" +def raise_func_escape(): + raise ValueError(expected_err_escape) + + +global_rref = None + + +def set_global_rref(rref): + global global_rref + global_rref = rref + + +def clear_global_rref(): + global global_rref + global_rref = None + + +def check_rref_confirmed(rref): + return rref.confirmed_by_owner() + + +def get_rref_debug_info(): + return _rref_context_get_debug_info() + + +def add_use_future_cb(to, x, y, z): + out = concurrent.futures.Future() + + def callback(fut): + out.set_result(fut.wait() + z) + + fut = rpc.rpc_async(to, torch.add, args=(x, y)) + fut.then(callback) + return out.result() + + +def get_events_from_profile(profile_rref): + return profile_rref.local_value().process_global_function_events + + +def add_use_future_set_result(to, x, y, z): + out = torch.futures.Future() + fut = rpc.rpc_async(to, torch.add, args=(x, y)) + fut.then(lambda fut : out.set_result(fut.wait() + z)) + return out.wait() + + +def add_use_future_nested_cb(to, x, y, z): + out = torch.futures.Future() + + def callback(fut1): + fut2 = rpc.rpc_async(to, torch.add, args=(fut1.wait(), z)) + fut2.then(lambda fut2 : out.set_result(fut2.wait())) + + fut1 = rpc.rpc_async(to, torch.add, args=(x, y)) + fut1.then(callback) + return out.wait() + + +def fail_on_fut(fut): + pass + + +@rpc.functions.async_execution +def async_raise_func(): + raise RuntimeError("Expected error") + + +@rpc.functions.async_execution +def async_wrong_type(): + return torch.zeros(2, 2) + + +@rpc.functions.async_execution +def async_add(to, x, y): + return rpc.rpc_async(to, torch.add, args=(x, y)) + + +def slow_add(x, y, device="cpu"): + time.sleep(1) + x = x.to(device) + y = y.to(device) + return torch.add(x, y).cpu() + + +@rpc.functions.async_execution +def slow_async_add(to, x, y, device="cpu"): + return rpc.rpc_async(to, slow_add, args=(x, y, device)) + + +@rpc.functions.async_execution +def async_add_with_future_ctor(to, x, y, z): + fut = torch.futures.Future() + rpc.rpc_async(to, torch.add, args=(x, y)).then( + lambda fut1: fut.set_result(fut1.wait() + z) + ) + return fut + + +@rpc.functions.async_execution +def async_add_chained(to, x, y, z): + return rpc.rpc_async(to, torch.add, args=(x, y)).then( + lambda fut: fut.wait() + z + ) + + +@rpc.functions.async_execution +def async_add_chained_multi(to, x, num, step): + fut = rpc.rpc_async(to, torch.add, args=(x, 0)) + for _ in range(num): + fut = fut.then(lambda fut: fut.wait() + step) + return fut + + +@rpc.functions.async_execution +def async_add_nested(to, x, y, z): + return rpc.rpc_async(to, async_add, args=(to, x, y)).then( + lambda fut: fut.wait() + z + ) + + +@rpc.functions.async_execution +def async_add_multi_fanout(to, x, num, step): + futs = [] + for i in range(num): + if i == 0: + futs.append(rpc.rpc_async(to, torch.add, args=(x, step))) + else: + futs.append(rpc.rpc_async(to, torch.add, args=(0, step))) + + # TODO: use torch.futures.collect_all + lock = Lock() + state = {"cnt": 0, "ret": torch.zeros_like(x)} + ret_future = torch.futures.Future() + + def inc_and_set(fut): + with lock: + state["cnt"] += 1 + state["ret"] += fut.wait() + if state["cnt"] >= len(futs): + ret_future.set_result(state["ret"]) + + for fut in futs: + fut.then(inc_and_set) + + return ret_future + + +@rpc.functions.async_execution +def async_cuda_sleep_and_set_to_one(t): + device = t.device + original_stream = torch.cuda.current_stream(device) + new_stream = torch.cuda.Stream(device) + new_stream.wait_stream(original_stream) + with torch.cuda.stream(new_stream): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + t.fill_(1) + fut = Future(devices=[device]) + fut.set_result(t) + return fut + + +@rpc.functions.async_execution +def async_cuda_nested_add(to, x, y, z): + def cb(fut): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + return fut.value() + z + + return rpc.rpc_async(to, torch.add, args=(x, y)).then(cb) + + +# A custom Python class that contains a tensor, needed to see if we correctly +# use the Python pickler to extract tensors from non-IValue-convertible types. +class TensorWrapper: + __slots__ = ("tensor", "lock", "event", "thread") + + def __init__(self, t): + self.tensor = t + # Add one non-picklable field, to ensure it's ignored/skipped. + self.lock = Lock() + self.event = torch.cuda.Event(enable_timing=True) + self.thread = threading.Thread() + self.thread.start() + + def increase(self, v): + with self.lock: + self.tensor += v + + def sum(self): + with self.lock: + self.event.record() + return self.tensor.sum() + + +class AsyncExecutionClass: + + @staticmethod + @rpc.functions.async_execution + def static_async_add(to, x, y, z): + return rpc.rpc_async(to, torch.add, args=(x, y)).then( + lambda fut: fut.wait() + z + ) + + @classmethod + @rpc.functions.async_execution + def class_async_add(cls, to, x, y, z): + ret_fut = torch.futures.Future() + rpc.rpc_async(to, torch.add, args=(x, y)).then( + lambda fut: ret_fut.set_result(fut.wait() + z) + ) + return ret_fut + + @rpc.functions.async_execution + def bound_async_add(self, to, x, y, z): + return rpc.rpc_async(to, torch.add, args=(x, y)).then( + lambda fut: fut.wait() + z + ) + + +def return_future(): + return torch.futures.Future() + + +class FooBackendOptions(rpc.RpcBackendOptions): + def __init__(self, init_method): + # Must call the __init__ of the superclass (and do so directly, + # without using super()) because... pybind. + rpc.RpcBackendOptions.__init__(self) + self.init_method = init_method + + +# load_tests from common_utils is used to automatically filter tests for +# sharding on sandcastle. This line silences flake warnings +load_tests = load_tests + + +class MyEmbeddingBagModel(torch.nn.Module): + def __init__(self, sparse): + super().__init__() + self.eb = torch.nn.EmbeddingBag( + 10, + 10, + sparse=sparse + ) + + def forward(self, x): + return self.eb(x) + + +class MyParameterServer: + def __init__(self, trainers): + self.lock = Lock() + self.trainers = trainers + self.iteration = 0 + self.updates = 0 + self.futures = [] + self.total = None + self.gradient = None + + @staticmethod + def get_gradient(rref): + return rref.local_value().gradient + + @staticmethod + @rpc.functions.async_execution + def average(rref, riteration, tensor): + self = rref.local_value() + fut = torch.futures.Future() + with self.lock: + if riteration > self.iteration: + self.iteration = riteration + self.updates = 0 + self.futures.clear() + self.futures.append(fut) + if self.total is None: + self.total = tensor + else: + self.total += tensor + self.updates += 1 + if self.trainers == self.updates: + self.gradient = self.total / float(self.trainers) + for fut in self.futures: + result = self.total / float(self.trainers) + fut.set_result(result) + return fut + + +class MyConvNetForMNIST(nn.Module): + def __init__(self, device): + super().__init__() + self.net = nn.Sequential( + nn.Conv2d(1, 16, 3, 1), + nn.ReLU(), + nn.Conv2d(16, 32, 3, 1), + nn.ReLU(), + nn.MaxPool2d(2), + nn.Flatten(1), + nn.Linear(4608, 128), + nn.ReLU(), + nn.Linear(128, 10), + ).to(device) + self.device = device + + def forward(self, x, is_rref=False): + x = x.to_here() if is_rref else x + with torch.cuda.stream(torch.cuda.current_stream(self.device)): + # intentionally adding delay to current CUDA stream + torch.cuda._sleep(10 * FIFTY_MIL_CYCLES) + return self.net(x) + + def __getstate__(self): + # return an empty dict to avoid inspecting the model contents on the + # owner + return {} + + +class RpcTestCommon: + def _run_func_in_mode(self, to, fn, mode, args=None, kwargs=None): + if mode == RPCExecMode.SYNC: + return rpc.rpc_sync(to, fn, args=args, kwargs=kwargs) + elif mode == RPCExecMode.ASYNC: + return rpc.rpc_async(to, fn, args=args, kwargs=kwargs).wait() + elif mode == RPCExecMode.REMOTE: + return rpc.remote(to, fn, args=args, kwargs=kwargs).to_here() + + def _self_py_udf_remote(self, worker_info, x, y, z): + rref = rpc.remote(worker_info, my_function, args=(x, y, z)) + self.assertEqual(rref.to_here(), x + y + z) + + def _self_remote_rref_as_rpc_arg(self, dst, x, y, z): + self_worker_info = rpc.get_worker_info() + rref = rpc.remote(self_worker_info, my_function, args=(x, y, z)) + fut = rpc.rpc_async(dst, add_rref_to_value, args=(rref, x)) + ret = rpc.rpc_sync(dst, add_rref_to_value, args=(rref, x + y)) + self.assertEqual(ret, x + y + z + x + y) + self.assertEqual(fut.wait(), x + y + z + x) + + def _self_remote_rref_as_remote_arg(self, dst, x, y, z): + self_worker_info = rpc.get_worker_info() + rref = rpc.remote(self_worker_info, my_function, args=(x, y, z)) + ret_rref = rpc.remote(dst, add_rref_to_value, args=(rref, x)) + self.assertEqual( + ret_rref.to_here(), x + y + z + x + ) + + def _world_size_one(self, a, b): + if self.rank == 0: + rpc.init_rpc( + name="me", + backend=self.rpc_backend, + rank=0, + world_size=1, + rpc_backend_options=self.rpc_backend_options, + ) + + def _rpc_sync(x, y): + expect = x * 2 + result = rpc.rpc_sync( + "me", + my_tensor_function, + args=(x, y) + ) + self.assertEqual(expect, result) + + def _rpc_async(x, y): + expect = x * 2 + result = rpc.rpc_async( + "me", + my_tensor_function, + args=(x, y) + ).wait() + self.assertEqual(expect, result) + + def _remote(x, y): + expect = x * 2 + result = rpc.remote( + "me", + my_tensor_function, + args=(x, y) + ).to_here() + self.assertEqual(expect, result) + + _rpc_sync(a, b) + _rpc_async(a, b) + _remote(a, b) + + rpc.shutdown() + + def _multi_rpc(self, sparse): + dst_rank = (self.rank + 1) % self.world_size + for i in range(20): + n = i + self.rank + 1 + if sparse: + x = build_sparse_tensor() * n + y = build_sparse_tensor() * n + else: + x = torch.ones(2, 2) + y = torch.ones(2, 2) + ret = rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(x, y), + ) + self.assertEqual(ret, x * 2) + + def _run_uneven_workload(self, f, x, num_repeat=30): + # worker0 drives and waits for worker1 and worker2 + # throughout the test. + if self.rank == 0: + self.assertTrue(self.world_size >= 3) + + # Phase 1: Only worker1 has workload. + dst = "worker1" + futs = [] + for _ in range(num_repeat): + fut = rpc.rpc_async(dst, f, args=(x,)) + futs.append(fut) + + for fut in torch.futures.collect_all(futs).wait(): + self.assertEqual(fut.wait(), 0) + + # Phase 2: Only worker2 has workload. + # If join is not correctly implemented, + # worker2 should be closed by now. + dst = "worker2" + futs = [] + for _ in range(num_repeat): + fut = rpc.rpc_async(dst, f, args=(x,)) + futs.append(fut) + + for val in torch.futures.wait_all(futs): + self.assertEqual(val, 0) + + def _wait_all_workers(self, f, x): + initialize_pg(self.file_init_method, self.rank, self.world_size) + rpc.init_rpc( + name="worker%d" % self.rank, + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + self._run_uneven_workload(f, x) + + # worker0 calls this at the end after waiting for RPC responses. + # worker1/2 calls this immediately and has some works after it. + # worker3 calls this immediately and has no more work. + rpc.api._wait_all_workers() + + # Wait before proceeding to shutdown to ensure worker0 RPCs make + # it through to other workers. + dist.barrier() + rpc.shutdown(graceful=False) + + def _wait_all_workers_twice(self, f, x): + initialize_pg(self.file_init_method, self.rank, self.world_size) + rpc.init_rpc( + name="worker%d" % self.rank, + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + self._run_uneven_workload(f, x) + + # worker0 calls this at the end after waiting for RPC responses. + # worker1/2 calls this immediately and has some works after it. + # worker3 calls this immediately and has no more work. + rpc.api._wait_all_workers() + rpc.api._wait_all_workers() + + # Wait before proceeding to shutdown to ensure worker0 RPCs make + # it through to other workers. + dist.barrier() + rpc.shutdown(graceful=False) + + def _nested_rpc(self, f, expected): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + f, + args=(worker_name(self.rank),), + ) + self.assertEqual(ret, expected) + + def _stress_test_rpc(self, f, repeat=1000, args=()): + n = self.rank + 1 + dst_rank = n % self.world_size + futs = [] + tik = time.time() + for _ in range(repeat): + fut = rpc.rpc_async(worker_name(dst_rank), f, args=args) + futs.append(fut) + + for val in torch.futures.wait_all(futs): + self.assertEqual(val, 0) + tok = time.time() + print( + f"Rank {self.rank} finished testing {repeat} times in {tok - tik} seconds." + ) + + def _builtin_remote_ret(self, x, y, expected): + n = self.rank + 1 + dst_rank = n % self.world_size + rref = rpc.remote( + worker_name(dst_rank), + torch.add, + args=(x, y), + ) + self.assertEqual(rref.to_here(), expected) + + def _builtin_remote_self(self, x, y, expected): + rref = rpc.remote( + worker_name(self.rank), + torch.add, + args=(x, y), + ) + self.assertEqual(rref.local_value(), expected) + + def _test_multi_remote_call(self, fn, sparse, args_fn=lambda x, y: (), kwargs_fn=lambda x, y: {}): + m = 10 + n = self.rank + 1 + dst_rank = n % self.world_size + rrefs = [] + expected = [] + for i in range(m): + n = n + i + rrefs.append( + rpc.remote( + worker_name(dst_rank), + fn, + args=args_fn(n, sparse), + kwargs=kwargs_fn(n, sparse), + ) + ) + expected.append(fn(*args_fn(n, sparse), **kwargs_fn(n, sparse))) + + for i in range(m): + self.assertEqual(rrefs[i].to_here(), expected[i]) + + def _py_rref_args(self, a, b, x, y, expected): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_a = rpc.remote( + worker_name(dst_rank), torch.add, args=(a, b) + ) + rref_b = rpc.remote( + worker_name(dst_rank), torch.add, args=(x, y) + ) + rref_c = rpc.remote( + worker_name(dst_rank), my_rref_function, args=(rref_a, rref_b) + ) + self.assertEqual(rref_c.to_here(), expected) + + def _py_rref_args_user_share(self, a, b, c, x, y, z, expected): + n = self.rank + 1 + owner_rank = n % self.world_size + user_rank = (n + 1) % self.world_size + rref_a = rpc.remote( + worker_name(owner_rank), my_function, args=(a, b, c) + ) + rref_b = rpc.remote( + worker_name(owner_rank), my_function, args=(x, y, z) + ) + rref_c = rpc.remote( + worker_name(user_rank), my_rref_function, args=(rref_a, rref_b) + ) + self.assertEqual(rref_c.to_here(), expected) + + def _py_rpc_rref_args(self, a, b, c, x, y, z, expected): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_a = rpc.remote( + worker_name(dst_rank), my_function, args=(a, b, c) + ) + rref_b = rpc.remote( + worker_name(dst_rank), my_function, args=(x, y, z) + ) + + c = rpc.rpc_sync( + worker_name(dst_rank), my_rref_function, args=(rref_a, rref_b) + ) + self.assertEqual(c, expected) + + def _nested_remote(self, f, expected): + n = self.rank + 1 + dst_rank1 = n % self.world_size + dst_rank2 = (n + 1) % self.world_size + + rref = rpc.remote( + worker_name(dst_rank1), + f, + args=(worker_name(dst_rank2),), + ) + self.assertEqual(rref.to_here(), expected) + + def _nested_rref(self, f, expected1, expected2): + n = self.rank + 1 + dst_rank1 = n % self.world_size + dst_rank2 = (n + 1) % self.world_size + rref_of_rrefs = rpc.remote( + worker_name(dst_rank1), + f, + args=(worker_name(dst_rank2),), + ) + + # Say C has 2 OwnerRRefs. + # B has 2 UserRRefs to those 2 OwnerRRefs, respectively. + # This call is effectively A asking B to share its 2 UserRRefs. + rrefs = rref_of_rrefs.to_here() + + self.assertEqual(len(rrefs), 2) + self.assertEqual(rrefs[0].to_here(), expected1) + self.assertEqual(rrefs[1].to_here(), expected2) + + def _nested_rref_stress(self, f, expected1, expected2): + n = self.rank + 1 + dst_rank1 = n % self.world_size + dst_rank2 = (n + 1) % self.world_size + all_rrefs = [ + rpc.remote( + worker_name(dst_rank1), + f, + args=(worker_name(dst_rank2),), + ) for _ in range(20) + ] + + for i in range(20): + rref_of_rrefs = all_rrefs[i] + rrefs = rref_of_rrefs.to_here() + self.assertEqual(len(rrefs), 2) + self.assertEqual(rrefs[0].to_here(), expected1) + self.assertEqual(rrefs[1].to_here(), expected2) + + def _trainer_func(self, rref, sparse): + m = MyEmbeddingBagModel(sparse=sparse) + loss_fn = nn.MSELoss() + for i in range(10): + outputs = m(torch.rand(10, 10).long()) + loss_fn(outputs, torch.rand(10, 10)).backward() + gradient = next(iter(m.parameters())).grad + fut = rref.rpc_async().average(rref, i, gradient) + gradient = fut.wait() + if gradient.is_sparse: + gradient = gradient.to_dense().double() + ps_gradient = rref.rpc_sync().get_gradient(rref) + if ps_gradient.is_sparse: + ps_gradient = ps_gradient.to_dense().double() + self.assertTrue(torch.equal(gradient, ps_gradient)) + + def _my_parameter_server(self, sparse): + ps_rref = RRef(MyParameterServer(self.world_size - 1)) + futures = [ + rpc.rpc_async( + worker_name((self.rank + index) % self.world_size), + self._trainer_func, + args=( + ps_rref, + sparse + ), + ) for index in range(1, self.world_size)] + torch.futures.wait_all(futures) + + def _test_cuda_future_extraction(self, wrapper, unwrapper, sparse_tensor): + # We check proper CUDA stream synchronization by adding to the tensor + # in one stream to get the expected value, and reading it from another stream. + future = Future(devices=["cuda:0"]) + with torch.cuda.device("cuda:0"): + stream = torch.cuda.Stream() + another_stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + if sparse_tensor: + tensor = build_sparse_tensor().to("cuda:0") + add_tensor = build_sparse_tensor().to("cuda:0") + expected_tensor = (tensor + add_tensor).coalesce() + else: + tensor = torch.zeros((100,), device="cuda:0") + add_tensor = torch.ones((100,), device="cuda:0") + expected_tensor = tensor + add_tensor + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + tensor += add_tensor + if sparse_tensor: + tensor = tensor.coalesce() + future.set_result(wrapper(tensor)) + with torch.cuda.stream(another_stream): + tensor = unwrapper(future.wait()) + if sparse_tensor: + self.assertTrue(torch.eq(tensor.indices(), expected_tensor.indices()).all().item()) + self.assertTrue(torch.eq(tensor.values(), expected_tensor.values()).all().item()) + self.assertEqual(tensor.size(), expected_tensor.size()) + else: + self.assertTrue(torch.eq(tensor, expected_tensor).all().item()) + + +class RpcTest(RpcAgentTestFixture, RpcTestCommon): + @dist_init + def test_worker_id(self): + n = self.rank + 1 + peer_rank = n % self.world_size + self_worker_info = rpc.get_worker_info() + peer_worker_info = rpc.get_worker_info(worker_name(peer_rank)) + + self.assertEqual(self_worker_info.name, worker_name(self.rank)) + self.assertEqual(peer_worker_info.name, worker_name(peer_rank)) + + with self.assertRaisesRegex(RuntimeError, "could not find destination"): + rpc.get_worker_info("WorkerUnknown") + + @dist_init + def test_get_worker_infos(self): + worker_infos = rpc.api._get_current_rpc_agent().get_worker_infos() + + worker_names = {worker_info.name for worker_info in worker_infos} + expected_worker_names = { + worker_name(rank) for rank in range(self.world_size) + } + self.assertEqual(worker_names, expected_worker_names) + + worker_ids = {worker_info.id for worker_info in worker_infos} + expected_worker_ids = set(range(self.world_size)) + self.assertEqual(worker_ids, expected_worker_ids) + + @dist_init + def test_self_add(self): + self_worker_info = rpc.get_worker_info() + fut = rpc.rpc_async(self_worker_info, torch.add, args=(torch.ones(2, 2), 1)) + ret = rpc.rpc_sync(self_worker_info, torch.add, args=(torch.ones(2, 2), 1)) + self.assertEqual(fut.wait(), torch.ones(2, 2) + 1) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + @dist_init + def test_send_to_rank(self): + dst_rank = (self.rank + 1) % self.world_size + + # Test dense tensor + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + ret = self._run_func_in_mode(dst_rank, torch.add, exec_mode, args=(torch.ones(2, 2), 1)) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + # Test invalid ranks + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + with self.assertRaises(RuntimeError): + self._run_func_in_mode(self.world_size + 1, torch.add, exec_mode, args=(torch.ones(2, 2), 1)) + + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + with self.assertRaises(RuntimeError): + self._run_func_in_mode(-1, torch.add, exec_mode, args=(torch.ones(2, 2), 1)) + + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + with self.assertRaises(ValueError): + self._run_func_in_mode(dst_rank + 0.5, torch.add, exec_mode, args=(torch.ones(2, 2), 1)) + + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + with self.assertRaises(ValueError): + self._run_func_in_mode(dst_rank - 0.5, torch.add, exec_mode, args=(torch.ones(2, 2), 1)) + + @dist_init + def test_self_py_udf_remote(self): + self._self_py_udf_remote( + rpc.get_worker_info(), + torch.ones(2, 2), + 1, + 3 + ) + + @dist_init + def test_self_remote_rref_as_rpc_arg(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._self_remote_rref_as_rpc_arg( + dst, + torch.ones(2, 2), + 1, + 3 + ) + + @dist_init + def test_self_remote_rref_as_self_rpc_arg(self): + self._self_remote_rref_as_rpc_arg( + rpc.get_worker_info(), + torch.ones(2, 2), + 1, + 3 + ) + + @dist_init + def test_self_remote_rref_as_remote_arg(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._self_remote_rref_as_remote_arg( + dst, + torch.ones(2, 2), + 1, + 3 + ) + + @dist_init + def test_self_remote_rref_as_self_remote_arg(self): + self._self_remote_rref_as_remote_arg( + rpc.get_worker_info(), + torch.ones(2, 2), + 1, + 3 + ) + + @dist_init + def test_rref_proxy_non_exist(self): + dst = worker_name((self.rank + 1) % self.world_size) + rref = rpc.remote(dst, my_function, args=(torch.ones(2, 2), 1, 3)) + msg = "has no attribute 'non_exist'" + with self.assertRaisesRegex(AttributeError, msg): + rref.rpc_sync().non_exist() + + with self.assertRaisesRegex(AttributeError, msg): + rref.rpc_async().non_exist().wait() + + with self.assertRaisesRegex(AttributeError, msg): + rref.remote().non_exist() + + def _test_rref_proxy_tensor(self, dst): + rref = rpc.remote(dst, my_function, args=(torch.ones(2, 2), 1, 3)) + + expected = torch.ones(2, 2) + 1 + 3 + self.assertEqual(expected.size(), rref.rpc_sync().size()) + self.assertEqual(expected + 1, rref.rpc_async().add(1).wait()) + self.assertEqual(expected.view(1, 4), rref.remote().view(1, 4).to_here()) + + @dist_init + def test_rref_proxy_tensor(self): + self._test_rref_proxy_tensor(worker_name((self.rank + 1) % self.world_size)) + + @dist_init + def test_rref_proxy_tensor_self(self): + self._test_rref_proxy_tensor(rpc.get_worker_info()) + + @dist_init + def test_rref_proxy_reuse(self): + rref = rpc.remote( + worker_name((self.rank + 1) % self.world_size), + my_function, + args=(torch.ones(2, 2), 1, 3) + ) + expected = torch.ones(2, 2) + 1 + 3 + + proxy_rpc_sync = rref.rpc_sync() + proxy_rpc_async = rref.rpc_async() + proxy_remote = rref.remote() + + self.assertEqual(expected.size(), proxy_rpc_sync.size()) + self.assertEqual(expected + 1, proxy_rpc_sync.add(1)) + self.assertEqual(expected.view(1, 4), proxy_rpc_sync.view(1, 4)) + + self.assertEqual(expected.size(), proxy_rpc_async.size().wait()) + self.assertEqual(expected + 3, proxy_rpc_async.add(3).wait()) + self.assertEqual(expected.view(4, 1), proxy_rpc_async.view(4, 1).wait()) + + self.assertEqual(expected.size(), proxy_remote.size().to_here()) + self.assertEqual(expected + 5, proxy_remote.add(5).to_here()) + self.assertEqual(expected.view(-1), proxy_remote.view(-1).to_here()) + + def _test_rref_proxy_class(self, dst): + rref = rpc.remote(dst, MyClass, args=(7,)) + expected = MyClass(7) + self.assertEqual(expected.get_value(), rref.rpc_sync().get_value()) + self.assertEqual(expected.get_value(), rref.rpc_async().get_value().wait()) + self.assertEqual(expected.get_value(), rref.remote().get_value().to_here()) + + expected.increment_value(3) + self.assertEqual(None, rref.rpc_sync().increment_value(1)) + self.assertEqual(None, rref.rpc_async().increment_value(1).wait()) + self.assertEqual(None, rref.remote().increment_value(1).to_here()) + + self.assertEqual(expected.get_value(), rref.rpc_sync().get_value()) + self.assertEqual(expected.get_value(), rref.rpc_async().get_value().wait()) + self.assertEqual(expected.get_value(), rref.remote().get_value().to_here()) + + self.assertEqual( + expected.my_instance_method(2), + rref.rpc_sync().my_instance_method(2) + ) + self.assertEqual( + expected.my_instance_method(3), + rref.rpc_async().my_instance_method(3).wait() + ) + self.assertEqual( + expected.my_instance_method(4), + rref.remote().my_instance_method(4).to_here() + ) + + self.assertEqual( + expected.my_static_method(9), + rref.rpc_sync().my_static_method(9) + ) + self.assertEqual( + expected.my_static_method(10), + rref.rpc_async().my_static_method(10).wait() + ) + self.assertEqual( + expected.my_static_method(11), + rref.remote().my_static_method(11).to_here() + ) + + self.assertEqual( + expected.my_class_method(2, torch.zeros(2, 2)), + rref.rpc_sync().my_class_method(2, torch.zeros(2, 2)) + ) + self.assertEqual( + expected.my_class_method(2, torch.ones(3, 3)), + rref.rpc_async().my_class_method(2, torch.ones(3, 3)).wait() + ) + self.assertEqual( + expected.my_class_method(2, torch.ones(4, 4)), + rref.remote().my_class_method(2, torch.ones(4, 4)).to_here() + ) + + @dist_init + def test_rref_proxy_class(self): + self._test_rref_proxy_class(worker_name((self.rank + 1) % self.world_size)) + + @dist_init + def test_rref_proxy_class_self(self): + self._test_rref_proxy_class(rpc.get_worker_info()) + + @mock.patch.object(torch.distributed.autograd, "_init") + @mock.patch.object(torch.distributed.rpc.api, "_set_and_start_rpc_agent") + @dist_init(setup_rpc=False) + def test_register_rpc_backend_and_set_and_start_rpc_backend( + self, mock_rpc_agent, mock_dist_autograd_init + ): + backend_name = "stub_backend" + + backend = rpc.backend_registry.register_backend( + backend_name, + _stub_construct_rpc_backend_options_handler, + _stub_init_rpc_backend_handler, + ) + + with self.assertRaisesRegex( + RuntimeError, "^RPC backend .+: already registered$" + ): + backend = rpc.backend_registry.register_backend( + backend_name, + _stub_construct_rpc_backend_options_handler, + _stub_init_rpc_backend_handler, + ) + + rpc.init_rpc( + name="worker1", + backend=backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + @dist_init(setup_rpc=False) + def test_duplicate_name(self): + with self.assertRaisesRegex(RuntimeError, "is not unique"): + store, _, _ = next( + torch.distributed.rendezvous( + self.init_method, rank=self.rank, world_size=self.world_size + ) + ) + rpc._init_rpc_backend( + backend=self.rpc_backend, + store=store, + name="duplicate_name", + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + @dist_init(setup_rpc=False) + def test_duplicate_name_2(self): + with self.assertRaisesRegex(RuntimeError, "is not unique"): + rpc.init_rpc( + name=worker_name(self.rank % (self.world_size - 1)), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + @dist_init(setup_rpc=False) + def test_reinit(self): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + initialize_pg(self.file_init_method, self.rank, self.world_size) + # Wait for all init to complete. + dist.barrier() + + # TODO: with TCP init, rank 0 raises Address already in use because + # rank 0 is the start daemon and the store is created before checking if + # RPC is already initialized in init_rpc. + if os.environ.get("RPC_INIT_WITH_TCP", None) == "1" and self.rank == 0: + expected_reinit_err = "Address already in use" + else: + expected_reinit_err = "is already initialized" + + with self.assertRaisesRegex(RuntimeError, expected_reinit_err): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + rpc.shutdown() + + @dist_init(setup_rpc=False) + def test_pg_init_no_rpc_init(self): + dist.init_process_group( + backend='gloo', + init_method=self.file_init_method, + rank=self.rank, + world_size=self.world_size) + + class MyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin = torch.nn.Linear(3, 4) + + def forward(self, x): + return self.lin(x) + + model = MyModel() + model.train() + model = torch.nn.parallel.DistributedDataParallel(model) + + with self.assertRaisesRegex(RuntimeError, 'Current RPC agent is not set! Did you initialize the RPC framework'): + params = [RRef(param) for param in model.parameters()] + + def test_world_size_one(self): + self._world_size_one( + torch.ones(2, 2), + torch.ones(2, 2) + ) + + @dist_init(setup_rpc=False) + def test_invalid_names(self): + + worker_id = 0 + with self.assertRaisesRegex(RuntimeError, "Worker name must match"): + WorkerInfo("abc*", worker_id) + + with self.assertRaisesRegex(RuntimeError, "Worker name must match"): + WorkerInfo(" ", worker_id) + + with self.assertRaisesRegex(RuntimeError, "must be non-empty"): + WorkerInfo("", worker_id) + + # If the number in the message does not match, it is likely that the + # value of MAX_NAME_LEN in RPC WorkerInfo has changed. + with self.assertRaisesRegex(RuntimeError, "shorter than 128"): + WorkerInfo("".join(["a" for i in range(500)]), worker_id) + + # Test that WorkerInfo can be pickled and sent in RPC call + @dist_init + def test_worker_info_pickle(self): + dst_rank = (self.rank + 1) % self.world_size + worker_info = rpc.api.get_worker_info() + ret = rpc.rpc_sync(worker_name(dst_rank), identity, args=(worker_info,)) + self.assertEqual(ret, worker_info) + + @dist_init + def test_add(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + self.assertEqual(ret, torch.ones(n, n) * 2) + + @staticmethod + def return_callee_id(): + return rpc.get_worker_info().id + + @dist_init + def test_int_callee(self): + dst_rank = (self.rank + 1) % self.world_size + ret = rpc.rpc_sync(dst_rank, RpcTest.return_callee_id) + self.assertEqual(ret, dst_rank) + + @dist_init + def test_add_with_id(self): + n = self.rank + 1 + dst_rank = n % self.world_size + workder_info = rpc.get_worker_info(worker_name(dst_rank)) + + ret = rpc.rpc_sync( + workder_info, torch.add, args=(torch.ones(n, n), torch.ones(n, n)) + ) + self.assertEqual(ret, torch.ones(n, n) * 2) + + @dist_init + def test_scalar_add(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), torch.add, args=(torch.ones(n, n), n) + ) + self.assertEqual(ret, (torch.ones(n, n) + n)) + + @dist_init + def test_async_add(self): + n = self.rank + 1 + dst_rank = n % self.world_size + fut = rpc.rpc_async( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + @dist_init + def test_nonzero(self): + n = self.rank + 1 + dst_rank = n % self.world_size + x = torch.ones(self.world_size, self.world_size) + x[self.rank][self.rank] = 0 + ret = rpc.rpc_sync(worker_name(dst_rank), torch.nonzero, args=(x,)) + self.assertEqual(ret, x.nonzero()) + + @dist_init + def test_multi_rpc(self): + self._multi_rpc(False) + + @dist_init + def test_future_wait_twice(self): + dst = worker_name((self.rank + 1) % self.world_size) + futs = [rpc.rpc_async(dst, raise_func) for _ in range(20)] + + with self.assertRaisesRegex(ValueError, "Expected error"): + torch.futures.wait_all(futs) + + for fut in futs: + with self.assertRaisesRegex(ValueError, "Expected error"): + fut.wait() + + @dist_init(setup_rpc=False) + def test_wait_all_workers_timeout(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + og_func = rpc.api._wait_all_workers + + def wait_all_workers_sleep(timeout): + rpc.api._all_gather(SlowPickleClass(0.5), timeout=timeout) + + rpc.api._wait_all_workers = wait_all_workers_sleep + + try: + with self.assertRaisesRegex(RuntimeError, ''): + rpc.shutdown(graceful=True, timeout=0.01) + finally: + rpc.api._wait_all_workers = og_func + dist.barrier() + + def test_wait_all_workers_dense(self): + self._wait_all_workers(heavy_rpc, torch.ones(100, 100)) + + def test_wait_all_workers_twice_dense(self): + self._wait_all_workers_twice(heavy_rpc, torch.ones(100, 100)) + + @dist_init + def test_all_gather(self): + info = rpc.get_worker_info() + results = rpc.api._all_gather(info.id) + expected = {} + for info in rpc._get_current_rpc_agent().get_worker_infos(): + expected[info.name] = info.id + + self.assertEqual(expected, results) + + @dist_init + def test_all_gather_timeout(self): + rpc._set_rpc_timeout(0.1) + + if self.rank == 0: + with self.assertRaisesRegex( + RuntimeError, + "timed out in _all_gather after 0\\.10 seconds" + ): + rpc.api._all_gather(SlowPickleClass(0.5)) + else: + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.api._all_gather(SlowPickleClass(0.5)) + + def _test_barrier_helper(self, info, names, multi_threaded=False): + names = sorted(names) + leader = names[0] + rpc.rpc_sync(leader, _reset_count) + if not multi_threaded and info.name == leader: + self.assertEqual(_rpc_barrier_count, 0) + rpc.api._barrier(names) + rpc.rpc_sync(leader, _increment_count) + rpc.api._barrier(names) + if not multi_threaded and info.name == leader: + self.assertEqual(_rpc_barrier_count, len(names)) + + @dist_init + def test_rpc_barrier_all(self): + # Test rpc barrier when called with full list of workers + info = rpc.get_worker_info() + all_worker_info = rpc._get_current_rpc_agent().get_worker_infos() + names = [worker.name for worker in all_worker_info] + self._test_barrier_helper(info, names) + + @dist_init + def test_rpc_barrier_subset(self): + # Test rpc barrier when processes are called with different subsets of the full list + info = rpc.get_worker_info() + all_worker_info = rpc._get_current_rpc_agent().get_worker_infos() + if info.id % 2: + names = [worker.name for worker in all_worker_info if worker.id % 2] + else: + names = [worker.name for worker in all_worker_info if not worker.id % 2] + self._test_barrier_helper(info, names) + + @dist_init + def test_rpc_barrier_partial_subset(self): + # Test rpc barrier when some processes are not involved in the barrier + info = rpc.get_worker_info() + all_worker_info = rpc._get_current_rpc_agent().get_worker_infos() + if info.id % 2: + names = [worker.name for worker in all_worker_info if worker.id % 2] + else: + names = [f"worker{info.id}"] + self._test_barrier_helper(info, names) + + @dist_init + def test_rpc_barrier_multithreaded(self): + # This tests validates the implementation of barrier when multiple threads call into it + # We only need to check that it does not hang in this case + info = rpc.get_worker_info() + all_worker_info = rpc._get_current_rpc_agent().get_worker_infos() + names = [worker.name for worker in all_worker_info] + threads = [] + for _ in range(3): + th = threading.Thread(target=self._test_barrier_helper, args=(info, names, True)) + threads.append(th) + th.start() + for th in threads: + th.join() + + @dist_init + def test_graceful_shutdown_with_uneven_workload(self): + """Test graceful termination.""" + self._run_uneven_workload(heavy_rpc, torch.ones(100, 100)) + + @dist_init(setup_rpc=False) + def test_shutdown_followed_by_rpc(self): + # Initialize RPC. + rpc.init_rpc( + name="worker%d" % self.rank, + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + self.assertEqual(ret, torch.ones(n, n) * 2) + rpc.shutdown() + + with self.assertRaisesRegex(RuntimeError, "^RPC has not been initialized"): + rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + + @dist_init + def test_expected_src(self): + dst_rank = (self.rank + 1) % self.world_size + expected_src_rank = (self.rank - 1) % self.world_size + rpc.rpc_sync(worker_name(dst_rank), set_value, args=(self.rank,)) + value = VALUE_FUTURE.result() + self.assertEqual(value, expected_src_rank) + + @dist_init + def test_py_built_in(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync(worker_name(dst_rank), min, args=(n, n + 1, n + 2)) + self.assertEqual(ret, min(n, n + 1, n + 2)) + + @dist_init + def test_py_user_defined(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + my_function, + kwargs={"a": n, "b": n + 1, "c": n + 2}, + ) + self.assertEqual(ret, my_function(n, n + 1, n + 2)) + + def test_build_rpc_profiling_key(self): + # Tests that the name that shows up as an Event in profiling RPCs has all + # the necessary information. + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + rpc_profiling_key = _build_rpc_profiling_key( + exec_mode, "foo", "worker0", "worker1" + ) + self.assertIn(exec_mode.value, rpc_profiling_key) + self.assertIn("foo", rpc_profiling_key) + self.assertIn("worker0", rpc_profiling_key) + self.assertIn("worker1", rpc_profiling_key) + + def check_profiling_info(self, self_worker_name, dst_worker_name, func, rpc_event, rpc_exec_mode): + self.assertTrue(self_worker_name in rpc_event.name) + self.assertTrue(dst_worker_name in rpc_event.name) + if isinstance(func, torch.jit.ScriptFunction): + self.assertTrue(torch._jit_internal._qualified_name(func) in rpc_event.name) + else: + self.assertTrue(func.__name__ in rpc_event.name) + self.assertTrue(rpc_exec_mode.value in rpc_event.name) + self.assertEqual(rpc_event.count, 1) + + @dist_init + def test_profiler_rpc_record_shapes(self): + if self.rank != 1: + return + dst = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst) + t1, t2 = torch.ones(100), torch.ones(100) + with _profile(record_shapes=True) as prof: + rpc.rpc_sync(dst_worker, torch.add, args=(t1, t2)) + + function_events = prof.function_events + remote_events = [event for event in function_events if event.is_remote] + remote_add_event = next( + event for event in remote_events if "aten::add" in event.name + ) + remote_add_input_shapes = remote_add_event.input_shapes + # Run profiler on equivalent local op and validate shapes are the same. + with _profile(record_shapes=True) as prof: + torch.add(t1, t2) + + local_function_events = prof.function_events + local_add_event = next( + event for event in local_function_events if "aten::add" in event.name + ) + local_add_input_shapes = local_add_event.input_shapes + self.assertEqual(remote_add_input_shapes, local_add_input_shapes) + + @dist_init + def test_profiler_rpc_memory(self): + if self.rank != 1: + return + dst = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst) + with _profile(profile_memory=True) as p: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + function_events = p.function_events + event_cpu_mem_usages = {event.cpu_memory_usage for event in function_events} + # if cpu_memory_usage was not propagated over the wire, this set would + # only contain 0 (indicates no memory being profiled) + self.assertNotEqual({0}, event_cpu_mem_usages) + # No memory profiled if profile_memory=False + with _profile(profile_memory=False) as p: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + function_events = p.function_events + event_cpu_mem_usages = {event.cpu_memory_usage for event in function_events} + self.assertEqual({0}, event_cpu_mem_usages) + + @dist_init + def test_profiler_export_trace(self): + if self.rank != 1: + return + dst = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst) + with _profile() as p: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + with TemporaryFileName() as fname: + path = fname + p.export_chrome_trace(path) + with open(path) as f: + trace = json.load(f) + event_names = [event['name'] for event in trace] + for expected_event_name in EXPECTED_REMOTE_EVENTS + [RPCExecMode.ASYNC.value]: + event_exists = any(expected_event_name in event_name for event_name in event_names) + self.assertTrue(event_exists) + + @dist_init + def test_profiler_rpc_key_names(self): + # tests that remote events are properly prefixed with the RPC profiling key. + if self.rank != 1: + return + + # Spawn multiple threads that send RPCs to ensure keys are correctly + # prefixed when there are multiple RPCs being created/in flight at the + # same time. + dst_ranks = [rank for rank in range(0, self.world_size) if rank != self.rank] + + def rpc_with_profiling(dst_worker): + with _profile() as prof: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + events = prof.function_events + remote_event_names = { + event.name: event for event in events if event.is_remote + } + rpc_profiling_key = _build_rpc_profiling_key( + RPCExecMode.ASYNC, + udf_with_torch_ops.__qualname__, + worker_name(self.rank), + dst_worker, + ) + + remote_event_name_set = set(EXPECTED_REMOTE_EVENTS) + for name, event in remote_event_names.items(): + # Ensure that we have the expected key as part of the remote + # event. + self.assertTrue(name.startswith(rpc_profiling_key)) + self.assertTrue(event.is_remote) + self.assertTrue(event.node_id == rpc.get_worker_info(dst_worker).id) + # Ensure that the remote event name also contains the operator. + operator_name_substr = name[len(rpc_profiling_key) :] + # Note: we don't assert that every remote event needs to be + # in the above set, the set is just a representative set of + # what we expect to see. The profiler can change and add more + # events, but we should always expect to see this representative + # set. + matching_event = { + remote_event_name + for remote_event_name in remote_event_name_set + if remote_event_name in operator_name_substr + } + remote_event_name_set -= matching_event + + # The set should be empty, otherwise its contained elements did + # not show up in the remote profiler output. + self.assertTrue( + remote_event_name_set == set(), + f"Expected {remote_event_name_set} to be included in remote profiler output.", + ) + + for dst in dst_ranks: + dst_worker = worker_name(dst) + num_parallel_rpcs = 2 + with concurrent.futures.ThreadPoolExecutor( + max_workers=num_parallel_rpcs + ) as executor: + futs = [ + executor.submit(rpc_with_profiling, dst_worker) + for _ in range(num_parallel_rpcs) + ] + # Wait for workers to finish test + for fut in futs: + fut.result() + + def _run_test_profiler_remote_events_profiled(self): + # Tests that we can successfully invoke the profiler on a remote node, + # and collect the remote events back in the local profiler. + if self.rank != 1: + return + + dst_ranks = [rank for rank in range(0, self.world_size) if rank != self.rank] + for dst in dst_ranks: + dst_worker = worker_name(dst) + with _profile() as prof: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + events = prof.function_events + + rpc_event = get_function_event(events, RPCExecMode.ASYNC.value) + self.check_profiling_info( + worker_name(self.rank), + dst_worker, + udf_with_torch_ops, + rpc_event, + RPCExecMode.ASYNC, + ) + + remote_events = {event.name: event for event in events if event.is_remote} + rpc_profiling_key = _build_rpc_profiling_key( + RPCExecMode.ASYNC, + udf_with_torch_ops.__qualname__, + worker_name(self.rank), + worker_name(dst), + ) + + for expected_remote_event_name in EXPECTED_REMOTE_EVENTS: + expected_key = rpc_profiling_key + REMOTE_OP_STR + expected_remote_event_name + self.assertTrue(expected_key in remote_events) + remote_event = remote_events[expected_key] + # Remote event should have a node ID corresponding to the worker + # it ran on. + self.assertEqual(remote_event.node_id, dst) + + # Validate order remote events show up in profiling output. + def convert_remote_to_local(event_name): + remote_op_key = rpc_profiling_key + REMOTE_OP_STR + return event_name[ + event_name.find(remote_op_key) + + len(remote_op_key) : + ] + + remote_events_list = [ + convert_remote_to_local(event.name) + for event in events + if convert_remote_to_local(event.name) in EXPECTED_REMOTE_EVENTS + ] + self.assertEqual( + set(remote_events_list), + set(EXPECTED_REMOTE_EVENTS), + f"Mismatch between profiled events: {set(remote_events_list)} and expected events: {set(EXPECTED_REMOTE_EVENTS)}", + ) + + @dist_init + def test_profiler_remote_events_profiled(self): + self._run_test_profiler_remote_events_profiled() + + @dist_init + def test_profiler_remote_events_profiled_single_threaded(self): + self._run_test_profiler_remote_events_profiled() + + def run_profiling_workload(self, dst): + fut = rpc.rpc_async( + worker_name(dst), + torch.mul, + args=( + torch.tensor(1.0, requires_grad=True), + torch.tensor(1.0, requires_grad=True), + ), + ) + fut.wait() + + def _run_rpc_profiling_async_function(self, device="cpu"): + if self.rank != 1: + return + + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + x = torch.ones(2) + y = torch.ones(2) + with _profile() as prof: + ret = rpc.rpc_async( + dst1, slow_async_add, args=(dst2, x, y, device), timeout=20 + ) + ret.wait() + + function_events = prof.function_events + # slow_async_add resulted in an RPC from dst1 -> dst2, so this should be + # recorded. + key_prefix = _build_rpc_profiling_key( + RPCExecMode.ASYNC, slow_async_add.__qualname__, worker_name(self.rank), dst1 + ) + + nested_rpc_key_prefix = _build_rpc_profiling_key( + RPCExecMode.ASYNC, slow_add.__qualname__, dst1, dst2 + ) + expected_key = key_prefix + REMOTE_OP_STR + nested_rpc_key_prefix + remote_events = [event for event in function_events if event.is_remote] + rpc_remote_event = [ + event for event in remote_events if event.name == expected_key + ] + self.assertEqual(1, len(rpc_remote_event)) + rpc_remote_event = rpc_remote_event[0] + self.assertEqual(rpc_remote_event.node_id, (self.rank + 1) % self.world_size) + # slow_async_add's RPC does an add on dst2, which should be reflected as well. + remote_add_key = ( + expected_key + REMOTE_OP_STR + torch.jit._builtins._find_builtin(torch.add) + ) + remote_add_event = [ + event for event in remote_events if event.name == remote_add_key + ] + self.assertEqual(1, len(remote_add_event)) + remote_add_event = remote_add_event[0] + # Validate that node_id is dst2. + self.assertEqual(remote_add_event.node_id, (self.rank + 2) % self.world_size) + + @dist_init + def test_rpc_profiling_async_function(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + self._run_rpc_profiling_async_function() + if torch.cuda.is_available(): + dist.barrier() + self._run_rpc_profiling_async_function(device="cuda:0") + + @dist_init + def test_rpc_profiling_async_function_single_threaded(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + self._run_rpc_profiling_async_function() + if torch.cuda.is_available(): + dist.barrier() + self._run_rpc_profiling_async_function(device="cuda:0") + + @dist_init + def test_rpc_profiling_remote_record_function(self): + # test that functions run over RPC with record_function show the expected + # profiled block. + if self.rank != 1: + return + dst_ranks = [i for i in range(self.world_size) if i != self.rank] + for dst_rank in dst_ranks: + dst_worker = worker_name(dst_rank) + with _profile() as prof: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=(-1, True)) + fut.wait() + + function_events = prof.function_events + record_function_remote_event = [ + evt for evt in function_events if "##forward##" in evt.name + ] + self.assertEqual(1, len(record_function_remote_event)) + record_function_remote_event = record_function_remote_event[0] + self.assertEqual(record_function_remote_event.node_id, dst_rank) + # cpu_children only returns direct children, so here we get all + # children recursively. + + def get_cpu_children(event): + if not event.cpu_children: + return [] + cpu_children = event.cpu_children + for e in event.cpu_children: + cpu_children.extend(get_cpu_children(e)) + return cpu_children + + remote_children = get_cpu_children(record_function_remote_event) + # Get local children and verify parity. + with _profile() as prof: + udf_with_torch_ops(-1, True) + + local_function_events = prof.function_events + local_record_function_event = next( + evt for evt in local_function_events if "##forward##" in evt.name + ) + local_children = get_cpu_children(local_record_function_event) + local_children_names = [ + evt.name for evt in local_children + ] + + REMOTE_OP_STR = "#remote_op: " + + def convert_remote_to_local(event_name): + remote_op_key = REMOTE_OP_STR + return event_name[ + event_name.find(remote_op_key) + len(remote_op_key) : + ] + + for evt in remote_children: + local_name = convert_remote_to_local(evt.name) + self.assertTrue(local_name in local_children_names) + + def validate_profiling_workload(self, dst, prof): + + def convert_remote_to_local(event_name): + return event_name[event_name.find(REMOTE_OP_STR) + len(REMOTE_OP_STR) :] + + events = prof.function_events + remote_events = { + convert_remote_to_local(event.name): event + for event in events + if event.is_remote + } + self.assertTrue("aten::mul" in remote_events) + remote_mul_event = remote_events["aten::mul"] + self.assertEqual(remote_mul_event.node_id, dst) + self.check_profiling_info( + worker_name(self.rank), + worker_name(dst), + torch.mul, + remote_mul_event, + RPCExecMode.ASYNC, + ) + + def _run_test_profiler_with_autograd_context(self): + dst = (self.rank + 1) % self.world_size + if self.rank == 1: + # Cases where we can double wrap messages with profiling information and autograd info. + with dist_autograd.context(): + with _profile() as prof: + self.run_profiling_workload(dst) + + self.validate_profiling_workload(dst, prof) + + # Ensure that flipped order of ctx managers results in events being + # recorded as expected. + with _profile() as prof: + with dist_autograd.context(): + self.run_profiling_workload(dst) + + self.validate_profiling_workload(dst, prof) + + @dist_init + def test_profiler_with_autograd_context_single_threaded(self): + self._run_test_profiler_with_autograd_context() + + @dist_init + def test_profiler_with_autograd_context(self): + self._run_test_profiler_with_autograd_context() + + def _profiler_test_with_rpc( + self, rpc_exec_mode, func, args, use_record_function=False, dst=None, kineto_profile=False + ): + dst = dst if dst is not None else (self.rank + 1) % self.world_size + + # only run profiler on rank 1. + p = _profile if not kineto_profile else torch.profiler.profile # kineto + if self.rank == 1: + with p() as prof: + record_function_ctx_mgr = ( + contextlib.nullcontext() + if not use_record_function + else torch.autograd.profiler.record_function( + "foo" + ) + ) + with record_function_ctx_mgr: + if rpc_exec_mode == RPCExecMode.SYNC: + rpc.rpc_sync(worker_name(dst), func, args=args) + elif rpc_exec_mode == RPCExecMode.ASYNC: + fut = rpc.rpc_async(worker_name(dst), func, args=args) + if kineto_profile: + # Ensure multiple async RPCs don't cause issues. + # Would have raised + # "RuntimeError: Cannot call + # RemoteProfilerManager::setCurrentKey when current + # key is already set." error if RPC profiling was + # not disabled properly for kineto. + fut2 = rpc.rpc_async(worker_name(dst), func, args=args) + fut2.wait() + fut.wait() + else: + self.assertTrue(rpc_exec_mode == RPCExecMode.REMOTE) + rref = rpc.remote(worker_name(dst), func, args=args) + rref.to_here() + # To avoid flakiness, wait for the RRef to be profiled. This + # means that we received the acknowledgement of successful + # creation on the owner and ran the callbacks responsible + # for recording the profiling event. + rref._get_profiling_future().wait() + + events = prof.function_events if not kineto_profile else prof.events() + if kineto_profile: + # RPC profiling is disabled so there should be no rpc related + # events. + with self.assertRaises(IndexError): + get_function_event(events, rpc_exec_mode.value) + + return + + rpc_event = get_function_event(events, rpc_exec_mode.value) + # verify Node ID for this rpc event. + self.assertEqual(rpc_event.node_id, self.rank) + # Ensure recording of remote events. + remote_events = {event for event in events if event.node_id == dst} - {rpc_event} + self.assertGreaterEqual(len(remote_events), 1) + for remote_event in remote_events: + self.assertEqual(remote_event.node_id, dst) + + if use_record_function: + scope_event = get_function_event(events, "foo") + # Since RPC call is within the scope, its CPU interval should be + # contained within foo's interval. + self.assertLessEqual(scope_event.time_range.start, rpc_event.time_range.start) + self.assertGreaterEqual(scope_event.time_range.end, rpc_event.time_range.end) + # the sender, dest worker, function run, and type of RPC should all + # be recorded. + self_worker_name = worker_name(self.rank) + dst_worker_name = worker_name(dst) + self.check_profiling_info(self_worker_name, dst_worker_name, func, rpc_event, rpc_exec_mode) + if use_record_function: + # verify order by ensuring that the outer context comes + # before the rpc event. + foo_event_ix = next(i for i, event in enumerate(events) if "foo" in event.name) + rpc_event_idx = next(i for i, event in enumerate(events) if rpc_exec_mode.value in event.name) + self.assertLess(foo_event_ix, rpc_event_idx) + + def _run_test_profiler_with_sync_rpc_udf(self): + self._profiler_test_with_rpc(RPCExecMode.SYNC, my_sleep_func, args=(1,)) + self._profiler_test_with_rpc(RPCExecMode.SYNC, my_sleep_func, args=(1,), + use_record_function=True) + + @dist_init + def test_profiler_with_sync_rpc_udf(self): + self._run_test_profiler_with_sync_rpc_udf() + + @dist_init + def test_profiler_with_sync_rpc_udf_single_threaded(self): + self._run_test_profiler_with_sync_rpc_udf() + + def _run_test_profiler_with_sync_rpc_builtin(self): + self._profiler_test_with_rpc( + RPCExecMode.SYNC, torch.mul, args=(torch.ones(1), torch.ones(1)) + ) + self._profiler_test_with_rpc( + RPCExecMode.SYNC, torch.mul, args=(torch.ones(1), torch.ones(1)), + use_record_function=True + ) + + @dist_init + def test_profiler_with_sync_rpc_builtin(self): + self._run_test_profiler_with_sync_rpc_builtin() + + @dist_init + def test_profiler_with_sync_rpc_builtin_single_threaded(self): + self._run_test_profiler_with_sync_rpc_builtin() + + def _run_test_profiler_with_async_rpc_udf(self): + self._profiler_test_with_rpc(RPCExecMode.ASYNC, my_sleep_func, args=(1,)) + self._profiler_test_with_rpc(RPCExecMode.ASYNC, my_sleep_func, args=(1,), + use_record_function=True) + # Test to ensure that kineto profiler enabled in RPC does not enable + # RPC profiling (it is unsupported) and does not result in issues. + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, my_sleep_func, args=(1,), kineto_profile=True + ) + + @dist_init + def test_profiler_with_async_rpc_udf(self): + self._run_test_profiler_with_async_rpc_udf() + + @dist_init + def test_profiler_with_async_rpc_udf_single_threaded(self): + self._run_test_profiler_with_async_rpc_udf() + + def _run_test_profiler_with_async_rpc_builtin(self): + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, torch.mul, args=(torch.ones(1), torch.ones(1)) + ) + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, torch.mul, args=(torch.ones(1), torch.ones(1)), + use_record_function=True + ) + + @dist_init + def test_profiler_with_async_rpc_builtin(self): + self._run_test_profiler_with_async_rpc_builtin() + + @dist_init + def test_profiler_with_async_rpc_builtin_single_threaded(self): + self._run_test_profiler_with_async_rpc_builtin() + + def _run_test_profiler_with_remote_udf(self): + self._profiler_test_with_rpc(RPCExecMode.REMOTE, my_sleep_func, args=(1,)) + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, my_sleep_func, args=(1,), use_record_function=True + ) + # test remote to self + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, my_sleep_func, args=(1,), dst=self.rank + ) + + @dist_init + def test_profiler_with_remote_udf(self): + self._run_test_profiler_with_remote_udf() + + @dist_init + def test_profiler_with_remote_udf_single_threaded(self): + self._run_test_profiler_with_remote_udf() + + def _run_test_profiler_with_remote_builtin(self): + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, torch.mul, args=(torch.ones(1), torch.ones(1)) + ) + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, torch.mul, args=(torch.ones(1), torch.ones(1)), + use_record_function=True + ) + # test remote to self + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, + torch.mul, + args=(torch.ones(1), torch.ones(1)), + dst=self.rank, + ) + + @dist_init + def test_profiler_with_remote_builtin(self): + self._run_test_profiler_with_remote_builtin() + + @dist_init + def test_profiler_with_remote_builtin_single_threaded(self): + self._run_test_profiler_with_remote_builtin() + + def _run_test_profiler_with_script_async_rpc(self): + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, my_script_func, args=(torch.tensor(1),) + ) + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, + my_script_func, + args=(torch.tensor(1),), + use_record_function=True, + ) + + @dist_init + def test_profiler_with_script_async_rpc(self): + self._run_test_profiler_with_script_async_rpc() + + @dist_init + def test_profiler_with_script_async_rpc_single_threaded(self): + self._run_test_profiler_with_script_async_rpc() + + def _run_test_profiler_with_script_sync_rpc(self): + self._profiler_test_with_rpc( + RPCExecMode.SYNC, my_script_func, args=(torch.tensor(1),) + ) + self._profiler_test_with_rpc( + RPCExecMode.SYNC, + my_script_func, + args=(torch.tensor(1),), + use_record_function=True, + ) + + @dist_init + def test_profiler_with_script_sync_rpc(self): + self._run_test_profiler_with_script_sync_rpc() + + @dist_init + def test_profiler_with_script_sync_rpc_single_threaded(self): + self._run_test_profiler_with_script_sync_rpc() + + def _run_test_profiler_with_script_remote_rpc(self): + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, my_script_func, args=(torch.tensor(1),) + ) + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, + my_script_func, + args=(torch.tensor(1),), + use_record_function=True, + ) + # test remote to self + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, my_script_func, args=(torch.tensor(1),), dst=self.rank + ) + + @dist_init + def test_profiler_with_script_remote_rpc(self): + self._run_test_profiler_with_script_remote_rpc() + + @dist_init + def test_profiler_with_script_remote_rpc_single_threaded(self): + self._run_test_profiler_with_script_remote_rpc() + + def _assert_top_level_events(self, process_global_events, expected_top_level_event_names): + top_level_event_names = [] + for thread_local_events in process_global_events: + # Get top-level events from all events happened on a thread. + last_end_time = 0 + for event in thread_local_events: + event_name = event.name + time_range = event.time_range + if time_range.start > last_end_time: + top_level_event_names.append(event_name) + last_end_time = time_range.end + top_level_event_names = sorted(top_level_event_names) + expected_top_level_event_names = sorted(expected_top_level_event_names) + self.assertEqual( + top_level_event_names, + expected_top_level_event_names, + f"Expected events {expected_top_level_event_names}, but got {top_level_event_names}", + ) + + @dist_init + def test_server_process_global_profiler(self): + if self.rank != 0: + return + + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = worker_name(dst_rank) + + x = torch.tensor(1) + y = torch.tensor(2) + + outer_profile_rref = rpc.remote(dst_worker_name, rpc._server_process_global_profile) + outer_profile_rref.rpc_sync().__enter__() + rpc.rpc_sync(dst_worker_name, torch.add, (x, y)) + inner_profile_rref = rpc.remote(dst_worker_name, rpc._server_process_global_profile) + inner_profile_rref.rpc_sync().__enter__() + rpc.rpc_sync(dst_worker_name, torch.sub, (x, y)) + inner_profile_rref.rpc_sync().__exit__(None, None, None) + outer_profile_rref.rpc_sync().__exit__(None, None, None) + + inner_events = rpc.rpc_sync(dst_worker_name, get_events_from_profile, (inner_profile_rref,)) + expected_inner_events = ['aten::sub'] + expected_outer_events = expected_inner_events + ['aten::add'] + + self._assert_top_level_events(inner_events, expected_inner_events) + outer_events = rpc.rpc_sync(dst_worker_name, get_events_from_profile, (outer_profile_rref,)) + self._assert_top_level_events(outer_events, expected_outer_events) + + inner_profile_rref.rpc_sync().key_averages() + outer_profile_rref.rpc_sync().key_averages() + + @dist_init + def test_async_record_function_double_end_callbacks(self): + num_sleep_seconds = 1 + if self.rank == 1: + # Validate that calling the function twice results in an error. + with _profile(): + with torch.autograd.profiler.record_function("foo") as rf: + fut = rpc.rpc_async( + worker_name(0), my_sleep_func, args=(num_sleep_seconds,) + ) + rf._call_end_callbacks_on_future(fut) + with self.assertRaisesRegex( + RuntimeError, "can only be called once." + ): + rf._call_end_callbacks_on_future(fut) + fut.wait() + + @dist_init + def test_async_record_function_legacy(self): + # Test the legacy _record_function ops work + # Note: These exist for backward compatibility with TorchScript + num_sleep_seconds = 1 + if self.rank == 1: + with _profile(): + try: + handle = torch.ops.profiler._record_function_enter("foo", None) + fut = rpc.rpc_async( + worker_name(0), my_sleep_func, args=(num_sleep_seconds,) + ) + torch.ops.profiler._call_end_callbacks_on_jit_fut(handle, fut) + finally: + torch.ops.profiler._record_function_exit(handle) + + fut.wait() + + @dist_init + def test_async_record_function_cbs_jit_call(self): + if self.rank == 1: + with _profile() as pf: + key = _build_rpc_profiling_key( + RPCExecMode.ASYNC, + torch._jit_internal._qualified_name(my_script_func), + "worker1", + "worker0", + ) + with torch.autograd.profiler.record_function(key) as rf: + fut = rpc.rpc_async( + worker_name(0), my_script_func, args=(torch.tensor(1),) + ) + # Intentionally calling record_function internals + fut = torch.ops.profiler._call_end_callbacks_on_jit_fut(rf.record, fut) + result = fut.wait() + # Validate that the profiling future returns the same value as the RPC + # future. + expected = torch.add(torch.tensor(1), torch.tensor(1)) + self.assertEqual(result, expected) + events = pf.function_events + rpc_event = get_function_event( + events, torch._jit_internal._qualified_name(my_script_func) + ) + self.assertTrue(torch._jit_internal._qualified_name(my_script_func) in rpc_event.name) + + @dist_init + def test_py_class_constructor(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync(worker_name(dst_rank), MyClass, args=(n,)) + self.assertEqual(ret.a, n) + + @dist_init + def test_py_class_instance_method(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), MyClass(2).my_instance_method, args=(n,) + ) + self.assertEqual(ret, MyClass(2).my_instance_method(n)) + + @dist_init + def test_py_class_method(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), MyClass.my_class_method, args=(n, n + 1) + ) + self.assertEqual(ret, MyClass.my_class_method(n, n + 1)) + + @dist_init + def test_py_class_static_method(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), MyClass.my_static_method, args=(n + 10,) + ) + self.assertEqual(ret, MyClass.my_static_method(n + 10)) + + @dist_init + def test_py_multi_async_call(self): + n = self.rank + 1 + dst_rank = n % self.world_size + dst_worker_info = rpc.get_worker_info(worker_name(dst_rank)) + fut1 = rpc.rpc_async(dst_worker_info, MyClass.my_static_method, args=(n + 10,)) + fut2 = rpc.rpc_async(dst_worker_info, min, args=(n, n + 1, n + 2)) + self.assertEqual(fut1.wait(), MyClass.my_static_method(n + 10)) + self.assertEqual(fut2.wait(), min(n, n + 1, n + 2)) + + @dist_init + def test_py_no_return_result(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync(worker_name(dst_rank), no_result) + self.assertEqual(ret, no_result()) + + @dist_init + def test_py_tensors(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + my_tensor_function, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + self.assertEqual(ret, my_tensor_function(torch.ones(n, n), torch.ones(n, n))) + + @dist_init + def test_py_tensors_multi_async_call(self): + futs = [] + n = self.rank + 1 + dst_rank = n % self.world_size + for i in range(100): + fut = rpc.rpc_async( + worker_name(dst_rank), + my_tensor_function, + args=(torch.ones(i, i), torch.ones(i, i)), + ) + futs.append(fut) + + j = 0 + for val in torch.futures.wait_all(futs): + self.assertEqual( + val, my_tensor_function(torch.ones(j, j), torch.ones(j, j)) + ) + j += 1 + + @dist_init + def test_py_tensors_in_container(self): + n = self.rank + 1 + dst_rank = n % self.world_size + a = [torch.ones(n, n), torch.ones(n, n)] + b = TensorClass(build_complex_tensors()) + c = {"foo": torch.ones(n, n), "bar": torch.ones(n, n)} + ret = rpc.rpc_sync( + worker_name(dst_rank), my_complex_tensor_function, args=(a, b, c) + ) + self.assertEqual(ret, my_complex_tensor_function(a, b, c)) + + @dist_init + def test_py_nested_pickle(self): + n = self.rank + 1 + dst_rank = n % self.world_size + + ret = rpc.rpc_sync( + worker_name(dst_rank), + run_nested_pickle, + args=(MyPickleClass(), torch.ones(2, 2)), + ) + + m = MyPickleClass() + m.set(my_tensor_function(torch.ones(2, 2), torch.ones(2, 2))) + self.assertEqual(ret, run_nested_pickle(m, torch.ones(2, 2))) + + @dist_init + def test_py_function_exception(self): + n = self.rank + 1 + dst_rank = n % self.world_size + with self.assertRaises(TypeError): + rpc.rpc_sync(worker_name(dst_rank), no_result, args=(10,)) + + @dist_init + def test_py_raise_in_user_func(self): + with captured_output() as (_, err): + # This barrier prevents a race condition where the main thread has + # not entered the context manager when the remote function runs. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + n = self.rank + 1 + dst_rank = n % self.world_size + fut = rpc.rpc_async(worker_name(dst_rank), raise_func) + with self.assertRaisesRegex(ValueError, expected_err): + fut.wait() + # This barrier prevents a race condition where the main thread exits + # context manager before the remote function has ran. + dist.barrier() + + # Validate that trainers log errors when running functions. + stderr_lines = err.getvalue() + self.assertTrue(expected_err in stderr_lines) + + @dist_init + def test_py_raise_in_user_func_escaped_str(self): + n = self.rank + 1 + dst_rank = n % self.world_size + fut = rpc.rpc_async(worker_name(dst_rank), raise_func_escape) + try: + fut.wait() + except ValueError as e: + msg = str(e) + # Ensure newlines are unescaped to provide a better repr of error. + self.assertEqual(msg, msg.encode("utf-8").decode("unicode_escape")) + else: + self.assertTrue(False, "expected raise_func_escape to raise ValueError.") + + @dist_init + def test_nested_rpc(self): + self._nested_rpc(nested_rpc, torch.ones(2, 2) + 1) + + @dist_init + def test_stress_light_rpc(self): + self._stress_test_rpc(light_rpc) + + @dist_init + def test_stress_heavy_rpc(self): + self._stress_test_rpc(heavy_rpc, repeat=20, args=(torch.ones(100, 100),)) + + @dist_init + def test_stress_heavy_rpc_torchscript(self): + self._stress_test_rpc(heavy_rpc_torchscript, repeat=20, args=(torch.ones(100, 100),)) + + @dist_init + def test_builtin_remote_ret(self): + self._builtin_remote_ret( + torch.ones(2, 2), + torch.ones(2, 2), + torch.ones(2, 2) * 2 + ) + + @dist_init + def test_builtin_remote_self(self): + self._builtin_remote_self( + torch.ones(2, 2), + torch.ones(2, 2), + torch.ones(2, 2) * 2 + ) + + @staticmethod + def _multi_args_fn(n, sparse=False): + if sparse: + return (build_sparse_tensor(), build_sparse_tensor()) + else: + return (torch.ones(n, n), torch.ones(n, n)) + + @dist_init + def test_multi_builtin_remote_ret(self): + self._test_multi_remote_call( + torch.add, False, + args_fn=RpcTest._multi_args_fn + ) + + @dist_init + def test_py_udf_remote(self): + n = self.rank + 1 + dst_rank = n % self.world_size + rref = rpc.remote( + worker_name(dst_rank), + my_function, + kwargs={"a": n, "b": n + 1, "c": n + 2}, + ) + self.assertEqual(rref.to_here(), my_function(n, n + 1, n + 2)) + + @staticmethod + def _multi_kwargs_fn(n, sparse=False): + if sparse: + return { + "a": build_sparse_tensor(), + "b": build_sparse_tensor(), + "c": build_sparse_tensor() + } + else: + return {"a": torch.ones(n, n), "b": torch.ones(n, n), "c": torch.ones(n, n)} + + @dist_init + def test_multi_py_udf_remote(self): + self._test_multi_remote_call( + my_function, + False, + kwargs_fn=RpcTest._multi_kwargs_fn + ) + + @dist_init + def test_py_rref_args(self): + self._py_rref_args( + torch.ones(2, 2), + 1, + torch.ones(2, 2), + 2, + torch.ones(2, 2) * 2 + 3) + + @dist_init + def test_py_rref_args_user_share(self): + self._py_rref_args_user_share( + torch.ones(2, 2), + 1, + 2, + torch.ones(2, 2), + 3, + 4, + torch.ones(2, 2) * 2 + 10 + ) + + @dist_init + def test_py_rpc_rref_args(self): + self._py_rpc_rref_args( + torch.ones(2, 2), + 1, + 2, + torch.ones(2, 2), + 3, + 4, + torch.ones(2, 2) * 2 + 10 + ) + + @dist_init + def test_nested_remote(self): + self._nested_remote( + nested_remote, + torch.ones(2, 2) + 3 + ) + + @dist_init + def test_nested_rref(self): + self._nested_rref( + nested_rref, + torch.ones(2, 2) + 1, + torch.ones(2, 2) + 2 + ) + + @dist_init + def test_nested_rref_stress(self): + self._nested_rref_stress( + nested_rref, + torch.ones(2, 2) + 1, + torch.ones(2, 2) + 2 + ) + + @dist_init + def test_multi_layer_nested_async_rpc(self): + # This test will exit right away, but there will be a chain of async + # RPCs. The termination algorithm should detect those messages properly. + # Otherwise, some peer could exit early, leaving others to timeout + # errors or connection closed errors. + ttl = 20 + n = self.rank + 1 + dst_rank = n % self.world_size + + multi_layer_nested_async_rpc(dst_rank, self.world_size, ttl) + + @dist_init + def test_remote_with_exception(self): + n = self.rank + 1 + dst_rank = n % self.world_size + # check ref to other workers + rref = rpc.remote(worker_name(dst_rank), raise_func) + with self.assertRaises(ValueError): + rref.to_here() + # check ref to itself + rref = rpc.remote(worker_name(self.rank), no_result, args=(10,)) + with self.assertRaises(TypeError): + rref.to_here() + + @dist_init + def test_rpc_return_rref(self): + n = self.rank + 1 + dst_rank1 = n % self.world_size + dst_rank2 = (n + 1) % self.world_size + rref = rpc.rpc_sync( + worker_name(dst_rank1), + rpc_return_rref, + args=(worker_name(dst_rank2),), + ) + self.assertEqual(rref.to_here(), torch.ones(2, 2) + 1) + + @dist_init + def test_rref_forward_chain(self): + ttl = 8 + n = self.rank + 1 + dst_rank = n % self.world_size + + rref = rpc.remote( + worker_name(dst_rank), torch.add, args=(torch.ones(n, n), 1) + ) + + ret_rref = rref_forward_chain(dst_rank, self.world_size, rref, ttl) + + for _ in range(ttl): + self.assertEqual(len(ret_rref), 1) + ret_rref = ret_rref[0].to_here() + + ret = ret_rref + self.assertEqual(ret, torch.add(torch.ones(n, n), 1)) + + @dist_init + def test_local_rref_no_fork(self): + local_rref = RRef(35) + self.assertEqual(local_rref.local_value(), 35) + + @dist_init + def test_local_value_not_on_owner(self): + # ensure that an error message is thrown if a user tries to call + # local_value() on a non-owning node. + next_rank = (self.rank + 1) % self.world_size + rref = rpc.remote( + worker_name(next_rank), torch.add, args=(torch.ones(1), torch.ones(1)) + ) + with self.assertRaisesRegex( + RuntimeError, ( + fr"For UserRRef\(rref_id=GloballyUniqueId\(created_on={self.rank}, local_id=0\), " + fr"fork_id=GloballyUniqueId\(created_on={self.rank}, local_id=1\)\), " + r"can't call localValue\(\) on user " + fr"WorkerInfo\(id={self.rank}, name={worker_name(self.rank)}\). " + fr"Call it on owner WorkerInfo\(id={next_rank}, name={worker_name(next_rank)}\)" + ) + ): + rref.local_value() + + @dist_init + def test_return_local_rrefs(self): + n = self.rank + 1 + dst_rank = n % self.world_size + + rref_list = rpc.rpc_sync( + worker_name(dst_rank), get_rref_list, args=([1, 2, 3],) + ) + + for rref in rref_list: + rpc.rpc_sync( + rref.owner(), + _call_method_on_rref, + args=(MyClass.increment_value, rref, 10), + ) + + rets = [ + rpc.rpc_sync( + rref.owner(), _call_method_on_rref, args=(MyClass.get_value, rref) + ) + for rref in rref_list + ] + + self.assertEqual(rets, [11, 12, 13]) + + @dist_init + def _test_rref_type(self, blocking): + + def launched_rpc(events): + expected_name = f"rpc_{RPCExecMode.ASYNC.value}#_rref_typeof_on_owner" + return any(e.name.startswith(expected_name) for e in events) + + dst = worker_name((self.rank + 1) % self.world_size) + rref = rpc.remote(dst, torch.add, args=(torch.ones(2), 1)) + + with _profile() as p: + t = rref._get_type(blocking=blocking) + if not blocking: + t = t.wait() + + self.assertTrue(launched_rpc(p.function_events)) + expected_type = type(torch.ones(2)) + self.assertEqual(t, expected_type) + + futs = [] + + def verify(fut): + self.assertEqual(fut.value(), expected_type) + + with _profile() as p: + for _ in range(10): + t = rref._get_type(blocking=blocking) + if not blocking: + futs.append(t) + t.add_done_callback(verify) + t = t.wait() + self.assertEqual(t, expected_type) + + if not blocking: + # Note that cached calls with blocking=False all return the same + # cached original future. + first_fut = futs[0] + for f in futs[1:]: + self.assertTrue(f is first_fut) + # Ensure we never launch another RPC, other than for the very + # first call. + self.assertFalse(launched_rpc(p.function_events)) + self.assertEqual(t, type(torch.ones(2))) + + rref = rpc.remote(dst, MyClass, args=(0,)) + rref_type = rref._get_type(blocking=blocking) + if not blocking: + rref_type = rref_type.wait() + self.assertEqual(rref_type, MyClass) + + def test_rref_type_blocking(self): + self._test_rref_type(blocking=True) + + def test_rref_type_non_blocking(self): + self._test_rref_type(blocking=False) + + @dist_init + def _test_rref_type_with_error(self, blocking): + dst = worker_name((self.rank + 1) % self.world_size) + # 10 ms timeout + rref = rpc.remote(dst, raise_func) + # Blocking: error raised inline + if blocking: + with self.assertRaisesRegex(ValueError, "Expected error"): + rref._get_type(blocking=blocking) + else: + # Non-blocking: Immediately return future, block on wait + fut = rref._get_type(blocking=blocking) + with self.assertRaisesRegex(ValueError, "Expected error"): + fut.wait() + + + def test_rref_type_with_error_blocking(self): + self._test_rref_type_with_error(blocking=True) + + def test_rref_type_with_error_non_blocking(self): + self._test_rref_type_with_error(blocking=False) + + @dist_init + def _test_rref_type_owner(self, blocking): + rref = RRef(torch.ones(2) + 1) + rref_type = rref._get_type(blocking=blocking) + if not blocking: + rref_type = rref_type.wait() + self.assertEqual(rref_type, type(torch.ones(2))) + + rref = RRef(MyClass(0)) + rref_type = rref._get_type(blocking=blocking) + if not blocking: + rref_type = rref_type.wait() + self.assertEqual(rref_type, MyClass) + + def test_rref_type_owner_blocking(self): + self._test_rref_type_owner(blocking=True) + + def test_rref_type_owner_non_blocking(self): + self._test_rref_type_owner(blocking=False) + + @staticmethod + def _slow_add(x, y): + time.sleep(1) + return x + y + + @dist_init + def test_rref_type_slow_init(self): + dst = worker_name((self.rank + 1) % self.world_size) + rref = rpc.remote(dst, RpcTest._slow_add, args=(torch.ones(2), 1)) + self.assertEqual(rref._get_type(), type(torch.ones(2))) + + @dist_init + def test_owner_equality(self): + a = RRef(40) + b = RRef(50) + + other_rank = (self.rank + 1) % self.world_size + other_a = rpc.remote( + worker_name(other_rank), torch.add, args=(torch.ones(1), 1) + ) + other_b = rpc.remote( + worker_name(other_rank), torch.add, args=(torch.ones(1), 1) + ) + other_a.to_here() # to ensure clean termination + other_b.to_here() + + self.assertNotEqual(a.owner(), 23) + self.assertEqual(other_a.owner(), other_b.owner()) + self.assertNotEqual(a.owner(), other_a.owner()) + self.assertEqual(other_a.owner(), other_a.owner()) + self.assertEqual(other_a.owner(), other_b.owner()) + self.assertEqual(a.owner(), a.owner()) + self.assertEqual(a.owner(), b.owner()) + self.assertEqual(a.owner(), rpc.get_worker_info()) + x = {} + x[a.owner()] = a + x[other_a.owner()] = other_a + self.assertEqual(x[a.owner()], a) + self.assertEqual(x[b.owner()], a) + self.assertEqual(x[other_a.owner()], other_a) + self.assertEqual(x[other_b.owner()], other_a) + self.assertEqual(len(x), 2) + + @dist_init + def test_pass_local_rrefs(self): + n = self.rank + 1 + dst_rank = n % self.world_size + dst_worker = worker_name(dst_rank) + + rref = RRef(40) + self.assertEqual( + rpc.rpc_sync(dst_worker, add_rref_to_value, args=(rref, 50)), 90 + ) + self.assertEqual( + rpc.rpc_async(dst_worker, add_rref_to_value, args=(rref, 50)).wait(), 90 + ) + self.assertEqual( + rpc.remote(dst_worker, add_rref_to_value, args=(rref, 50)).to_here(), 90 + ) + + @dist_init + def test_remote_same_worker(self): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_a = rpc.remote( + worker_name(dst_rank), torch.add, args=(torch.ones(n, n), 2) + ) + rref_b = rpc.remote( + worker_name(dst_rank), torch.add, args=(torch.ones(n, n), 1) + ) + rref_c = rpc.remote( + worker_name(dst_rank), my_rref_function, args=(rref_a, rref_b) + ) + self.assertEqual(rref_c.to_here(), torch.ones(n, n) + 4) + + @dist_init(setup_rpc=True) + def test_call_method_on_rref(self): + """ + Tests that it is possible to call an instance method on a remote object + by using rref.owner() as destination of the call. + """ + vals = [10, 2, 5, 7] + dst_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst_rank) + + # creates a remote object + rref = rpc.remote(dst_worker, MyClass, args=(vals[0],)) + + # modifies state of the remote object + rpc.rpc_sync( + rref.owner(), + _call_method_on_rref, + args=(MyClass.increment_value, rref, vals[1]), + ) + rpc.rpc_async( + rref.owner(), + _call_method_on_rref, + args=(MyClass.increment_value, rref, vals[2]), + ).wait() + rpc.remote( + rref.owner(), + _call_method_on_rref, + args=(MyClass.increment_value, rref, vals[3]), + ).to_here() + + # queries state of the remote object + result = rpc.rpc_sync( + dst_worker, _call_method_on_rref, args=(MyClass.get_value, rref) + ) + + self.assertEqual(result, sum(vals)) + + # Notice `rpc.api.shutdown()` accesses + # `_delete_all_user_and_unforked_owner_rrefs` through + # `torch.distributed.rpc.api`, so patching + # `torch.distributed.rpc._delete_all_user_and_unforked_owner_rrefs` will + # not help. + @mock.patch.object(torch.distributed.rpc.api, "_delete_all_user_and_unforked_owner_rrefs") + def _test_rref_leak(self, _mock_delete_all_user_and_unforked_owner_rrefs, ignore_leak): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + initialize_pg(self.file_init_method, self.rank, self.world_size) + # Wait for all init to complete. + dist.barrier() + + rref = rpc.remote( # noqa: F841 + worker_name((self.rank + 1) % self.world_size), + torch.add, + args=(torch.ones(2, 2), 1), + ) + + import torch.distributed.rpc.api as api + + if ignore_leak: + api._ignore_rref_leak = True + rpc.shutdown(graceful=True) + else: + api._ignore_rref_leak = False + with self.assertRaisesRegex(RuntimeError, "Leaking RRef"): + rpc.shutdown(graceful=True) + + @dist_init(setup_rpc=False) + def test_rref_leak(self): + self._test_rref_leak(ignore_leak=False) + + @dist_init(setup_rpc=False) + def test_ignore_rref_leak(self): + self._test_rref_leak(ignore_leak=True) + + @dist_init + def test_rref_str(self): + rref1 = RRef(self.rank) + id_class = "GloballyUniqueId" + self.assertEqual( + f"OwnerRRef({id_class}(created_on={self.rank}, local_id=0))", rref1.__str__() + ) + + dst_rank = (self.rank + 1) % self.world_size + rref2 = rpc.remote( + worker_name(dst_rank), torch.add, args=(torch.ones(2, 2), 1) + ) + self.assertEqual( + rref2.__str__(), + f"UserRRef(RRefId = {id_class}(created_on={self.rank}, local_id=1), " + f"ForkId = {id_class}(created_on={self.rank}, local_id=2))", + ) + + @dist_init + def test_rref_get_future(self): + # Tests that we can obtain the future corresponding to the creation of + # the RRef on remote end + if self.rank == 0: + # Builtin + rref = rpc.remote(worker_name(1), torch.add, args=(1, 1)) + rref.to_here() + fut = rref._get_future() + self.assertIsInstance(fut, torch._C.Future) + + # UDF + rref = rpc.remote(worker_name(1), foo_add, args=()) + rref.to_here() + fut = rref._get_future() + self.assertIsInstance(fut, torch._C.Future) + + # Script + rref = rpc.remote(worker_name(1), my_script_func, args=(torch.tensor(1), )) + rref.to_here() + fut = rref._get_future() + self.assertIsInstance(fut, torch._C.Future) + + + @dist_init + def test_rref_context_debug_info(self): + # This test checks local states that are modified by remote workers. + # This means that we would need barrier before and after every check. + # The barrier before the check makes sure that all previous states are + # cleared globally, the barrier after ensures that no following states + # change gets into the current check. + initialize_pg(self.file_init_method, self.rank, self.world_size) + + # Check 1: local RRef does not update owners_ map or add a pending user. + ################################################# + + rref1 = RRef(self.rank) + + # don't need a barrier here as local RRef is handled by this thread + info = _rref_context_get_debug_info() + self.assertIn("num_owner_rrefs", info) + self.assertIn("num_pending_users", info) + # RRef on local value is not added to context until shared across RPC + self.assertEqual(0, int(info["num_owner_rrefs"])) + self.assertEqual(0, int(info["num_pending_users"])) + # barrier after the check 1 + dist.barrier() + + # Check 2: Sharing RRef as an arg should update owners_ map + ########################################################### + + dst_rank = (self.rank + 1) % self.world_size + rpc.rpc_sync(worker_name(dst_rank), set_global_rref, args=(rref1,)) + + # barrier before check 2 + wait_until_pending_futures_and_users_flushed() + dist.barrier() + + info = _rref_context_get_debug_info() + self.assertIn("num_owner_rrefs", info) + self.assertEqual(1, int(info["num_owner_rrefs"])) + # no pending users since the fork is finished + self.assertEqual(0, int(info["num_pending_users"])) + # barrier after check 2 + dist.barrier() + + # clear states for check 2 + rpc.rpc_sync(worker_name(dst_rank), clear_global_rref) + + # Wait for owner rref to be cleared. + while int(info["num_owner_rrefs"]) != 0: + info = _rref_context_get_debug_info() + time.sleep(0.1) + dist.barrier() + + # Check 3: rpc.remote call should update owners_ map + #################################################### + rref2 = rpc.remote( + worker_name(dst_rank), torch.add, args=(torch.ones(2, 2), 1) + ) + rref3 = rpc.remote( + worker_name(dst_rank), torch.add, args=(torch.ones(2, 2), 1) + ) + rref2.to_here() + rref3.to_here() + + # barrier before check 3 + wait_until_pending_futures_and_users_flushed() + dist.barrier() + + info = _rref_context_get_debug_info() + self.assertIn("num_owner_rrefs", info) + self.assertEqual(2, int(info["num_owner_rrefs"])) + # no pending users since the fork is finished + self.assertEqual(0, int(info["num_pending_users"])) + + # barrier after check 3 + dist.barrier() + + @dist_init + def test_disable_gil_profiling(self): + # test that rpc.enable_gil_profiling(false) will result in + # GIL wait time not being recorded. + + # GIL profiling should be disabled by default. + dst_rank = (self.rank + 1) % self.world_size + rpc.rpc_sync( + worker_name(dst_rank), torch.add, args=(torch.ones(1), torch.ones(1)) + ) + info = rpc.api._get_current_rpc_agent().get_debug_info() + self.assertRaises(KeyError, lambda: info["agent.gil_average_wait_time_us"]) + rpc.enable_gil_profiling(True) + rpc.rpc_sync( + worker_name(dst_rank), torch.add, args=(torch.ones(1), torch.ones(1)) + ) + info = rpc.api._get_current_rpc_agent().get_debug_info() + self.assertIn("agent.gil_average_wait_time_us", info) + + @dist_init(setup_rpc=False) + def test_local_shutdown(self): + # test that we can start RPC and then immediately locally shutdown + # without sending any messages. + rpc.init_rpc( + name="worker%d" % self.rank, + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + # pass in graceful=False to ensure that we don't wait for other workers. + rpc.shutdown(graceful=False) + + @dist_init + def test_debug_info(self): + # only test keys in this test case. Values should be covered by + # individual module debug info tests + import torch.distributed.autograd as dist_autograd + + info = _get_debug_info() + rref_info = _rref_context_get_debug_info() + agent_info = rpc.api._get_current_rpc_agent().get_debug_info() + autograd_info = dist_autograd._get_debug_info() + common_keys = rref_info.keys() & agent_info.keys() & autograd_info.keys() + self.assertEqual(0, len(common_keys)) + expected = {} + expected.update(rref_info) + expected.update(agent_info) + expected.update(autograd_info) + # NB: Key ordering is only preserved in python 3.6+. So here, we + # manually check keys are equal. + for key in expected.keys(): + self.assertIn(key, info.keys()) + + for key in info.keys(): + self.assertIn(key, expected.keys()) + + @dist_init(setup_rpc=False) + @skip_but_pass_in_sandcastle_if( + IS_MACOS, + "Test is flaky on MacOS since libuv error handling is not as robust as TCP", + ) + def test_handle_send_exceptions(self): + # test that if a callee node has gone down, we raise an appropriate + # exception instead of just crashing. + rpc.init_rpc( + name="worker%d" % self.rank, + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + rpc._set_rpc_timeout(10) + # This barrier is needed to ensure that some workers do not exit before + # others have been brought up. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + if self.rank == 1: + dst_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst_rank) + # allow destination worker to exit without joining + error_str = self.get_shutdown_error_regex() + wait_until_node_failure(dst_rank, error_str) + fut = rpc.rpc_async(dst_worker, torch.add, args=(torch.ones(1), 3)) + # Shutdown sequence is not very well defined and as a result + # we can see any of the error messages defined in get_shutdown_error_regex. + with self.assertRaisesRegex(RuntimeError, error_str): + fut.wait() + # exit all workers non-gracefully. + rpc.shutdown(graceful=False) + + @dist_init + def test_deadlock(self): + # this test is copied from https://github.com/pytorch/pytorch/issues/45089 + if self.rank == 1: + dst1 = worker_name((self.rank + 1) % self.world_size) + x = torch.ones(2) + y = torch.ones(2) + rpc.rpc_async(dst1, RpcTest._slow_add, args=(x, y), timeout=15).wait() + + dist_initialized = dist.is_initialized() + if not dist_initialized: + dist.init_process_group( + backend="gloo", + init_method=self.file_init_method, + rank=self.rank, + world_size=self.world_size, + ) + + @dist_init(setup_rpc=False) + def test_local_shutdown_with_rpc(self): + # test that we can start RPC, send RPCs, and then run local shutdown. + rpc.init_rpc( + name="worker%d" % self.rank, + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + n = self.rank + 1 + dst_rank = n % self.world_size + rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + # A barrier is needed to ensure that all RPCs are processed. + # Otherwise, some RPCs can timeout since the receiving end + # has terminated. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + # pass in graceful=False to ensure that we don't wait for other workers. + rpc.shutdown(graceful=False) + + @dist_init(setup_rpc=False) + def test_set_and_get_default_rpc_timeout(self): + timeout = 0.5 + + # A new `RpcBackendOptions` is constructed + # when accessing `self.rpc_backend_options`. + rpc_backend_options = self.rpc_backend_options + rpc_backend_options.rpc_timeout = timeout + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + set_timeout = rpc.get_rpc_timeout() + self.assertEqual(timeout, set_timeout) + rpc.shutdown() + + @dist_init + def test_default_timeout_used(self): + """ + Tests that if no timeout is passed into rpc_async and rpc_sync, then the + default timeout is used. + """ + dst_rank = (self.rank + 1) % self.world_size + rpc._set_rpc_timeout(0.001) # 1 ms + # futures should time out and be marked with an exception indicating it as such. + futs = [ + rpc.rpc_async(worker_name(dst_rank), my_sleep_func, args=()) + for _ in range(10) + ] + expected_error = self.get_timeout_error_regex() + for fut in futs: + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # ensure that if a new timeout is set old futures don't time out but new ones do. + rpc._set_rpc_timeout(200) # 200 seconds + # create a longstanding RPC. + fut1 = rpc.rpc_async(worker_name(dst_rank), my_sleep_func, args=(1,)) + # now, set a short timeout. + rpc._set_rpc_timeout(0.001) + # fut2 should time out, fut1 should not. + fut2 = rpc.rpc_async(worker_name(dst_rank), my_sleep_func, args=(1,)) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut2.wait() + fut1.wait() + + # Zero timeout means infinity, so future should run to completion. + rpc._set_rpc_timeout(0) + rpc.rpc_async(worker_name(dst_rank), my_sleep_func, args=()).wait() + + # reset to default timeout so shutdown messages can process cleanly. + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + @dist_init + def test_rpc_timeouts(self): + # TODO: enable timeouts for rpc.remote/RRef (https://github.com/pytorch/pytorch/issues/33803) + dst_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst_rank) + timeout = 0.1 # 100 ms + expected_error = self.get_timeout_error_regex() + # Test async UDF + fut = rpc.rpc_async(dst_worker, my_sleep_func, args=(1,), timeout=timeout) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure run to completion if there is no timeout and we use the default + # RPC timeout. + rpc.rpc_async(dst_worker, my_sleep_func, args=(1,)).wait() + + # Test sync UDF + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,), timeout=timeout) + + # Ensure run to completion if there is no timeout and we use the default + # RPC timeout. + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,)) + + # If we set a default timeout for RPCs, it should be respected, though + # still overridden if we pass in a different timeout to the APIs. + rpc._set_rpc_timeout(0.001) + fut = rpc.rpc_async(dst_worker, my_sleep_func, args=(1,)) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,)) + + # The RPCs should run to completion since we override the timeout. + rpc.rpc_async(dst_worker, my_sleep_func, args=(1,), timeout=5).wait() + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,), timeout=5) + # Passing in a zero timeout should ensure that the RPC won't time out. + rpc.rpc_async(dst_worker, my_sleep_func, args=(1,), timeout=0).wait() + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,), timeout=0) + # Reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + def test_dist_init_decorator(self): + @dist_init(setup_rpc=False) + def test_func(self): + return "expected result" + + self.assertEqual(test_func(self), "expected result") + + @dist_init + def test_func(self): + return "expected result" + + self.assertEqual(test_func(self), "expected result") + + def test_use_rpc_pickler(self): + class TestPickler: + pass + + test_pickler = TestPickler() + with _use_rpc_pickler(test_pickler): + self.assertTrue(torch.distributed.rpc.api._default_pickler is test_pickler) + self.assertTrue( + torch.distributed.rpc.api._default_pickler is _internal_rpc_pickler + ) + + @dist_init + def test_wait_all(self): + with _wait_all(): + self.assertTrue(_thread_local_var.future_list == []) + dst = worker_name((self.rank + 1) % self.world_size) + fut = rpc.rpc_async(dst, torch.add, (torch.ones(2, 2), 1)) + self.assertTrue(len(_thread_local_var.future_list) == 1) + self.assertTrue(isinstance(_thread_local_var.future_list[0], torch._C.Future)) + self.assertTrue(fut.done()) + self.assertEqual(fut.wait(), torch.ones(2, 2) + 1) + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_wait_all_multiple_call(self): + with _wait_all(): + self.assertTrue(_thread_local_var.future_list == []) + dst = worker_name((self.rank + 1) % self.world_size) + for i in range(20): + fut = rpc.rpc_async(dst, torch.add, (torch.ones(i, i), 1)) + res = rpc.rpc_sync(dst, torch.add, (torch.ones(i, i), 1)) + self.assertEqual(res, torch.ones(i, i) + 1) + self.assertEqual(fut.wait(), torch.ones(i, i) + 1) + self.assertTrue(len(_thread_local_var.future_list) == 20) + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_wait_all_timeout(self): + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + with _wait_all(): + self.assertTrue(_thread_local_var.future_list == []) + dst = worker_name((self.rank + 1) % self.world_size) + timeout = 0.1 # 100 ms + rpc.rpc_async(dst, my_sleep_func, args=(1,), timeout=timeout) + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_wait_all_raise_in_user_func(self): + with self.assertRaises(ValueError): + with _wait_all(): + self.assertTrue(_thread_local_var.future_list == []) + dst = worker_name((self.rank + 1) % self.world_size) + rpc.rpc_async(dst, raise_func) + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_wait_all_raise_in_body(self): + with self.assertRaises(ValueError): + with _wait_all(): + raise_func() + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_custom_exception_throw_during_reconstruction(self): + """ + Test that we still throw info about the remote side exception even when + we cannot recreate it on client side. + """ + initialize_pg(self.file_init_method, self.rank, self.world_size) + if self.rank != 0: + exc_caught = False + dst = worker_name(0) + try: + rpc.rpc_sync(dst, custom_raise_func, args=()) + except RuntimeError as e: + exc_caught = True + msg = str(e) + print(f"Got msg {msg}") + self.assertTrue("Original exception on remote side was" in msg) + self.assertTrue("CustomException" in msg) + except BaseException as e: + raise RuntimeError( + f"Failure - expected RuntimeError, got {e}" + ) from e + finally: + self.assertTrue(exc_caught) + + dist.barrier() + + + timed_out_rpc_event = None + + @staticmethod + def timed_out_rpc(): + RpcTest.timed_out_rpc_event.wait() + + @dist_init + def test_wait_all_exit_early_python(self): + # Initialize the event in the subprocess. + RpcTest.timed_out_rpc_event = Event() + + # Wait for all processes to initialize event. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + + dst = worker_name((self.rank + 1) % self.world_size) + fut1 = rpc.rpc_async(dst, RpcTest.timed_out_rpc) + fut2 = rpc.rpc_async(dst, raise_func) + fut3 = rpc.rpc_async(dst, raise_func) + + # We should receive the error from fut2 + with self.assertRaisesRegex(ValueError, expected_err): + torch.futures.wait_all([fut1, fut2, fut3]) + + # Unblock RPC thread for fut1 + RpcTest.timed_out_rpc_event.set() + + @dist_init + def test_wait_all_exit_early_builtin(self): + # Initialize the event in the subprocess. + RpcTest.timed_out_rpc_event = Event() + + # Wait for all processes to initialize event. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + + dst = worker_name((self.rank + 1) % self.world_size) + fut1 = rpc.rpc_async(dst, RpcTest.timed_out_rpc) + fut2 = rpc.rpc_async(dst, torch.add, args=(torch.rand(10), torch.rand(5))) + fut3 = rpc.rpc_async(dst, torch.add, args=(torch.rand(10), torch.rand(5))) + + # We should receive the error from fut2 + with self.assertRaisesRegex(RuntimeError, "size of tensor"): + torch.futures.wait_all([fut1, fut2, fut3]) + + # Unblock RPC thread for fut1 + RpcTest.timed_out_rpc_event.set() + + @dist_init + def test_wait_all_exit_early_script_function(self): + # Initialize the event in the subprocess. + RpcTest.timed_out_rpc_event = Event() + + # Wait for all processes to initialize event. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + + dst = worker_name((self.rank + 1) % self.world_size) + fut1 = rpc.rpc_async(dst, RpcTest.timed_out_rpc) + fut2 = rpc.rpc_async(dst, raise_func_script, args=(expected_err,)) + fut3 = rpc.rpc_async(dst, raise_func_script, args=(expected_err,)) + + # We should receive the error from fut2 + with self.assertRaisesRegex(RuntimeError, expected_err): + torch.futures.wait_all([fut1, fut2, fut3]) + + # Unblock RPC thread for fut1 + RpcTest.timed_out_rpc_event.set() + + + @dist_init + def test_function_not_on_callee(self): + # test that if a function does not exist on a callee, we don't crash, + # instead we get an AttributeError indicating that the func does not exist. + this_module = sys.modules[__name__] + caller_worker = "worker0" + callee_worker = "worker1" + + if self.rank == 1: + # Use delattr to remove the binding of a func on this nodes + delattr(this_module, "foo_add") + # notify remote end that we have removed it. + rpc.rpc_sync(caller_worker, set_value, args=(self.rank,)) + + if self.rank == 0: + # func exists on caller, but not callee. + # wait for remote end to remove the binding of foo_add func. + wait_for_value_future() + # Ensure that we have the attribute on this module. Otherwise, the test could fail due to a caller-side pickling error. + self.assertTrue(hasattr(this_module, "foo_add")) + with self.assertRaisesRegex( + RuntimeError, "RPC pickler does not serialize" + ): + rpc.rpc_sync(callee_worker, foo_add, args=()) + + @dist_init + def test_non_garbage_collected_user_rref_due_to_local_circular_dependency(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + a = MyClass(1) + b = MyClass(2) + + # This is to make Python not garbage collect a and b. + a.other = b + b.other = a + + n = self.rank + a.rref = rpc.remote( + dst_worker_name, + torch.add, + args=(torch.ones(n, n), 2) + ) + + @dist_init(setup_rpc=False) + def test_use_rref_after_shutdown(self): + rpc.init_rpc( + name="worker%d" % self.rank, + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + n = self.rank + 1 + dst_rank = n % self.world_size + rref = rpc.remote( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + # pass in graceful=True to ensure that local UserRRefs are deleted. + rpc.shutdown(graceful=True) + + with self.assertRaisesRegex( + RuntimeError, "Cannot call to_here\\(\\) on it after deletion." + ): + rref.to_here() + + with self.assertRaisesRegex( + RuntimeError, "Cannot call fork an UserRRef after deletion." + ): + import torch.distributed.rpc.internal as internal + internal.serialize(rref) + + @staticmethod + def _return_gpu_tensor(): + return torch.rand(3, 3).cuda(0) + + @staticmethod + def _return_gpu_tensor_list(): + return [torch.rand(3, 3).cuda(0), torch.rand(3, 3).cuda(1)] + + @staticmethod + def _gpu_tensor_list_arg(tensor_list): + return torch.rand(3, 3) + + def _create_rref(self): + owner_rank = (self.rank + 2) % self.world_size + return rpc.remote( + worker_name(owner_rank), + torch.add, + args=(torch.zeros(2, 2), 1) + ) + + @dist_init + def test_user_rrefs_confirmed(self): + dst_rank = (self.rank + 1) % self.world_size + rref = self._create_rref() + ret = rpc.rpc_sync( + worker_name(dst_rank), + check_rref_confirmed, + args=(rref,) + ) + self.assertEqual(ret, True) + + @dist_init + def test_user_rrefs_confirmed_remote(self): + dst_rank = (self.rank + 1) % self.world_size + rref = self._create_rref() + ret_rref = rpc.remote( + worker_name(dst_rank), + check_rref_confirmed, + args=(rref,) + ) + self.assertEqual(ret_rref.to_here(), True) + + @dist_init + def test_rref_py_pickle_not_supported(self): + local_rref = RRef(35) + with TemporaryFileName() as fname: + with self.assertRaisesRegex(RuntimeError, "Can not pickle rref in python pickler"): + torch.save(local_rref, fname) + + @dist_init + def test_remote_throw(self): + rref = rpc.remote(worker_name((self.rank + 1) % self.world_size), + raise_or_inc, + args=(torch.ones(2),)) + with self.assertRaisesRegex(Exception, ".*Expected error.*"): + rref.to_here() + + @dist_init + def test_non_cont_tensors(self): + if self.rank == 0: + # Create a non-contiguous tensor. + t = torch.rand(5, 5) + t_view = t.narrow(1, 2, 2) + self.assertFalse(t_view.is_contiguous()) + t_cont = t_view.contiguous() + self.assertTrue(t_cont.is_contiguous()) + self.assertEqual(t_view, t_cont) + + # Send non-cont tensor over RPC. + next_rank = (self.rank + 1) % self.world_size + t_ret = rpc.rpc_sync(worker_name(next_rank), non_cont_test, args=(t_view, t_cont)) + + # Verify the returned tensor. + self.assertEqual(t_view, t_ret) + self.assertFalse(t_ret.is_contiguous()) + + @dist_init + def test_callback_simple(self): + set_by_cb = concurrent.futures.Future() + n = self.rank + 1 + + def callback(fut): + ret = fut.wait() + self.assertEqual(ret, torch.ones(n, n) * 2) + set_by_cb.set_result(ret.clone() + 1) + + fut = rpc.rpc_async( + worker_name(n % self.world_size), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)) + ) + + fut.then(callback) + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + self.assertEqual(set_by_cb.result(), torch.ones(n, n) * 2 + 1) + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + @dist_init + def test_callback_wrong_arg_num(self): + n = self.rank + 1 + + fut = rpc.rpc_async( + worker_name(n % self.world_size), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)) + ) + + cb_fut = fut.then(my_function) + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + with self.assertRaisesRegex( + RuntimeError, + "my\\_function\\(\\) missing 2 required positional arguments" + ): + cb_fut.wait() + + @dist_init + def test_callback_wrong_arg_type(self): + dst = worker_name((self.rank + 1) % self.world_size) + + fut0 = rpc.rpc_async(dst, torch.add, args=(torch.ones(2, 2), 1)) + fut1 = fut0.then(lambda x: x + 1) + + with self.assertRaisesRegex( + RuntimeError, + "unsupported operand type\\(s\\) for \\+" + ): + fut1.wait() + + @dist_init + def test_callback_multi(self): + num_cbs = 10 + n = self.rank + 1 + + def callback(idx, fut): + ret = fut.wait() + self.assertEqual(ret, torch.ones(n, n) * 2) + return ret + idx + + fut = rpc.rpc_async( + worker_name(n % self.world_size), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)) + ) + + cb_futs = [fut.then(partial(callback, idx)) for idx in range(num_cbs)] + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + for idx in range(num_cbs): + self.assertEqual( + cb_futs[idx].wait(), + torch.ones(n, n) * 2 + idx + ) + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + @dist_init + def test_callback_chain(self): + n = self.rank + 1 + + def callback(fut): + return fut.wait() + 1 + + fut = rpc.rpc_async( + worker_name(n % self.world_size), + torch.add, + args=(torch.ones(n, n), 1) + ) + + num_cbs = 20 + for _ in range(num_cbs): + fut = fut.then(callback) + + self.assertEqual(fut.wait(), torch.ones(n, n) + 1 + num_cbs) + + @dist_init + def test_callback_in_rpc(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + ret = rpc.rpc_sync( + dst1, + add_use_future_cb, + args=(dst2, torch.ones(2, 2), 1, 2) + ) + self.assertEqual(ret, torch.ones(2, 2) + 1 + 2) + + @dist_init + def test_callback_with_ret(self): + dst = worker_name((self.rank + 1) % self.world_size) + + def callback(fut0): + fut2 = rpc.rpc_async( + dst, + torch.add, + args=(fut0.wait(), 1) + ).then(lambda fut1: fut1.wait() + 1) + + return fut2.wait() + + fut3 = rpc.rpc_async( + dst, + torch.add, + args=(torch.ones(2, 2), 1) + ).then(callback) + + self.assertEqual(fut3.wait(), torch.ones(2, 2) + 3) + + @dist_init + def test_callback_with_error(self): + dst = worker_name((self.rank + 1) % self.world_size) + + def callback(fut0): + with self.assertRaisesRegex(ValueError, "Expected error"): + fut0.wait() + raise RuntimeError("Another expected error") + + fut1 = rpc.rpc_async(dst, raise_func).then(callback) + with self.assertRaisesRegex(RuntimeError, "Another expected error"): + fut1.wait() + + @dist_init + def test_callback_none(self): + dst = worker_name((self.rank + 1) % self.world_size) + with self.assertRaisesRegex( + TypeError, + "incompatible function arguments." + ): + rpc.rpc_async(dst, raise_func).then(None) + + @dist_init + def test_add_done_callback(self): + set_by_cb = False + n = self.rank + 1 + + def callback(fut): + nonlocal set_by_cb + fut.wait() + set_by_cb = True + + fut = rpc.rpc_async( + worker_name(n % self.world_size), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)) + ) + + fut.add_done_callback(callback) + fut_then = fut.then(lambda _: True) + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + # We have no guarantee that the add_done_callback fn will execute before the test finishes. + # Adding a 'then' callback that runs afterwards to guarantee we wait for the first callback + fut_then.wait() + self.assertTrue(set_by_cb) + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + @dist_init + def test_mark_future_twice(self): + fut = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + torch.add, + args=(torch.zeros(2, 2), 1) + ) + self.assertEqual(fut.wait(), torch.zeros(2, 2) + 1) + with self.assertRaisesRegex( + RuntimeError, + "Future can only be marked completed once" + ): + fut.set_result(1) + + @dist_init + def test_pickle_future(self): + fut = torch.futures.Future() + errMsg = "Can not pickle torch.futures.Future" + + dst = worker_name((self.rank + 1) % self.world_size) + with TemporaryFileName(): + with self.assertRaisesRegex(RuntimeError, errMsg): + rpc.rpc_sync(dst, fail_on_fut, args=(fut,)) + + with TemporaryFileName(): + with self.assertRaisesRegex(RuntimeError, errMsg): + rpc.rpc_async(dst, fail_on_fut, args=(fut,)) + + with TemporaryFileName(): + with self.assertRaisesRegex(RuntimeError, errMsg): + rpc.remote(dst, fail_on_fut, args=(fut,)) + + @dist_init + def test_future_done(self): + dst = worker_name((self.rank + 1) % self.world_size) + fut = rpc.rpc_async(dst, torch.add, args=(torch.zeros(2), 1)) + fut.wait() + self.assertTrue(fut.done()) + + @dist_init + def test_future_done_exception(self): + dst = worker_name((self.rank + 1) % self.world_size) + fut = rpc.rpc_async(dst, raise_func) + with self.assertRaisesRegex(ValueError, "Expected error"): + fut.wait() + self.assertTrue(fut.done()) + + def _test_future_cb(self, func): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + ret = rpc.rpc_sync( + dst1, + func, + args=(dst2, torch.ones(2, 2), 1, 2) + ) + self.assertEqual(ret, torch.ones(2, 2) + 1 + 2) + + @dist_init + def test_future_in_rpc(self): + self._test_future_cb(add_use_future_set_result) + + @dist_init + def test_future_nested_callback(self): + self._test_future_cb(add_use_future_nested_cb) + + def _test_async_function_raise(self, mode): + with self.assertRaisesRegex(RuntimeError, "Expected error"): + self._run_func_in_mode( + worker_name((self.rank + 1) % self.world_size), + async_raise_func, + mode + ) + + @dist_init + def test_async_function_raise(self): + self._test_async_function_raise(RPCExecMode.SYNC) + + @dist_init + def test_async_function_raise_async(self): + self._test_async_function_raise(RPCExecMode.ASYNC) + + @dist_init + def test_async_function_raise_remote(self): + self._test_async_function_raise(RPCExecMode.REMOTE) + + def _test_async_function_wrong_return_type(self, mode): + errMsg = ( + "Functions decorated with @rpc\\.async_function must return a " + "torch\\.futures\\.Future object," + ) + with self.assertRaisesRegex(RuntimeError, errMsg): + self._run_func_in_mode( + worker_name((self.rank + 1) % self.world_size), + async_wrong_type, + mode + ) + + @dist_init + def test_async_function_wrong_return_type(self): + self._test_async_function_wrong_return_type(RPCExecMode.SYNC) + + @dist_init + def test_async_function_wrong_return_type_async(self): + self._test_async_function_wrong_return_type(RPCExecMode.ASYNC) + + @dist_init + def test_async_function_wrong_return_type_remote(self): + self._test_async_function_wrong_return_type(RPCExecMode.REMOTE) + + @dist_init + def test_async_function_simple(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + ret = rpc.rpc_sync(dst1, async_add, args=(dst2, torch.ones(2, 2), 1)) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + def _test_async_function(self, fn, mode=RPCExecMode.SYNC): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + args = (dst2, torch.ones(2, 2), 1, 2) + ret = self._run_func_in_mode(dst1, fn, mode, args=args) + self.assertEqual(ret, torch.ones(2, 2) + 3) + + @dist_init + def test_async_function_with_future_ctor(self): + self._test_async_function(async_add_with_future_ctor) + + @dist_init + def test_async_function_with_future_ctor_remote(self): + self._test_async_function( + async_add_with_future_ctor, + RPCExecMode.REMOTE + ) + + @dist_init + def test_async_function_chained(self): + self._test_async_function(async_add_chained) + + @dist_init + def test_async_function_chained_remote(self): + self._test_async_function(async_add_chained, RPCExecMode.REMOTE) + + @dist_init + def test_async_function_nested(self): + self._test_async_function(async_add_nested) + + @dist_init + def test_async_function_nested_remote(self): + self._test_async_function(async_add_nested, RPCExecMode.REMOTE) + + @dist_init + def test_async_static_method(self): + self._test_async_function(AsyncExecutionClass.static_async_add) + + @dist_init + def test_async_static_method_remote(self): + self._test_async_function( + AsyncExecutionClass.static_async_add, + RPCExecMode.REMOTE + ) + + @dist_init + def test_async_class_method(self): + self._test_async_function(AsyncExecutionClass.class_async_add) + + @dist_init + def test_async_class_method_remote(self): + self._test_async_function( + AsyncExecutionClass.class_async_add, + RPCExecMode.REMOTE + ) + + def _test_test_async_class_rref_proxy(self, mode=RPCExecMode.SYNC): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + rref = rpc.remote(dst1, AsyncExecutionClass) + + x = torch.ones(2, 2) + y = torch.ones(2, 2) + 1 + if mode == RPCExecMode.SYNC: + ret = rref.rpc_sync().static_async_add(dst2, x, x, y) + ret += rref.rpc_sync().class_async_add(dst2, x, x, y) + ret += rref.rpc_sync().bound_async_add(dst2, x, x, y) + elif mode == RPCExecMode.ASYNC: + ret = rref.rpc_async().static_async_add(dst2, x, x, y).wait() + ret += rref.rpc_async().class_async_add(dst2, x, x, y).wait() + ret += rref.rpc_async().bound_async_add(dst2, x, x, y).wait() + elif mode == RPCExecMode.REMOTE: + ret = rref.remote().static_async_add(dst2, x, x, y).to_here() + ret += rref.remote().class_async_add(dst2, x, x, y).to_here() + ret += rref.remote().bound_async_add(dst2, x, x, y).to_here() + + self.assertEqual(ret, 3 * 4 * x) + + @dist_init + def test_async_class_rref_proxy(self): + self._test_test_async_class_rref_proxy() + + @dist_init + def test_async_class_rref_proxy_async(self): + self._test_test_async_class_rref_proxy(mode=RPCExecMode.ASYNC) + + @dist_init + def test_async_class_rref_proxy_remote(self): + self._test_test_async_class_rref_proxy(mode=RPCExecMode.REMOTE) + + def _test_async_function_multi(self, fn, mode=RPCExecMode.SYNC): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + num = 20 + step = 3 + args = (dst2, torch.ones(2, 2), num, step) + ret = self._run_func_in_mode(dst1, fn, mode, args=args) + self.assertEqual(ret, torch.ones(2, 2) + num * step) + + @dist_init + def test_async_function_multi_chained(self): + self._test_async_function_multi(async_add_chained_multi) + + @dist_init + def test_async_function_multi_chained_async(self): + self._test_async_function_multi( + async_add_chained_multi, + RPCExecMode.ASYNC + ) + + @dist_init + def test_async_function_multi_chained_remote(self): + self._test_async_function_multi( + async_add_chained_multi, + RPCExecMode.REMOTE + ) + + @dist_init + def test_async_function_multi_fanout(self): + self._test_async_function_multi(async_add_multi_fanout) + + @dist_init + def test_async_function_multi_fanout_async(self): + self._test_async_function_multi( + async_add_multi_fanout, + RPCExecMode.ASYNC + ) + + @dist_init + def test_async_function_multi_fanout_remote(self): + self._test_async_function_multi( + async_add_multi_fanout, + RPCExecMode.REMOTE + ) + + def _test_return_future(self, mode): + with self.assertRaisesRegex( + RuntimeError, + "Can not pickle torch.futures.Future" + ): + self._run_func_in_mode( + worker_name((self.rank + 1) % self.world_size), + return_future, + mode + ) + + @dist_init + def test_return_future(self): + self._test_return_future(RPCExecMode.SYNC) + + @dist_init + def test_return_future_async(self): + self._test_return_future(RPCExecMode.ASYNC) + + @dist_init + def test_return_future_remote(self): + self._test_return_future(RPCExecMode.REMOTE) + + @dist_init + def test_rref_timeout(self): + # This test is similar to ones in FaultyProcessGroupTest, but is meant to be + # run with other backends besides ProcessGroup. + if self.rank != 0: + return + + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + # 10 ms timeout + rref = rpc.remote(dst_worker, my_sleep_func, args=(2, ), timeout=0.01) + # Future corresponding to the remote creation should time out. + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rref._get_future().wait() + # Call to ensure pending callbacks are run. + wait_until_pending_futures_and_users_flushed() + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref.to_here() + + wait_until_owners_and_forks_on_rank(1, 1, rank=1) + + @dist_init(setup_rpc=False) + @skip_but_pass_in_sandcastle_if( + os.environ.get("RPC_INIT_WITH_TCP", None) == "1", + "init_pg_then_rpc does not work with TCP init, see https://github.com/pytorch/pytorch/issues/41614." + ) + def test_init_pg_then_rpc(self): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + # Test RPC. + next_rank = (self.rank + 1) % self.world_size + ret = rpc.rpc_sync(worker_name(next_rank), torch.add, args=(torch.ones(2, 2), 1)) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + # Test PG + dist.barrier() + + rpc.shutdown() + + @dist_init(setup_rpc=False) + @skip_but_pass_in_sandcastle_if( + os.environ.get("RPC_INIT_WITH_TCP", None) == "1", + "init_rpc_then_pg does not work with TCP init, see https://github.com/pytorch/pytorch/issues/41614." + ) + def test_init_rpc_then_pg(self): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + + # Test RPC. + next_rank = (self.rank + 1) % self.world_size + ret = rpc.rpc_sync(worker_name(next_rank), torch.add, args=(torch.ones(2, 2), 1)) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + # Test PG + dist.barrier() + + rpc.shutdown() + + @dist_init + def test_wait_all_with_exception(self): + dst = worker_name((self.rank + 1) % self.world_size) + futs = [rpc.rpc_async(dst, raise_func) for _ in range(10)] + + with self.assertRaisesRegex(ValueError, "Expected error"): + torch.futures.wait_all(futs) + + @dist_init + def test_wait_all_with_partial_exception(self): + dst = worker_name((self.rank + 1) % self.world_size) + futs = [rpc.rpc_async(dst, torch.add, args=(torch.ones(2), 1)) for _ in range(10)] + + futs.append(rpc.rpc_async(dst, raise_func)) + + with self.assertRaisesRegex(ValueError, "Expected error"): + torch.futures.wait_all(futs) + + @dist_init(setup_rpc=False) + @skip_but_pass_in_sandcastle_if( + os.environ.get("RPC_INIT_WITH_TCP", None) == "1", + "Test does not work with TCP init, see https://github.com/pytorch/pytorch/issues/46491", + ) + def test_init_rpc_twice(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + rpc.shutdown() + + # Wait for all init to complete. + dist.barrier() + + # Use a different file name for the next initialization + new_backend_options = self.rpc_backend_options + new_backend_options.init_method += "init_2" + + # Ensure rpc initialization works again. + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=new_backend_options, + ) + + # Verify RPCs work after re-init. + dst = worker_name((self.rank + 1) % self.world_size) + rpc.rpc_sync(dst, torch.add, args=(torch.ones(2, 2), 1)) + rpc.rpc_sync(dst, foo_add, args=()) + + rpc.shutdown() + + def test_wrong_types(self): + with self.assertRaisesRegex( + TypeError, + "Argument backend must be a member of BackendType", + ): + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + backend="TENSORPIPE", + ) + + with self.assertRaisesRegex( + TypeError, + "Argument rpc_backend_options must be an instance of RpcBackendOptions", + ): + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + backend=self.rpc_backend, + rpc_backend_options={"init_method": self.init_method} + ) + + def test_cannot_infer_backend_from_options(self): + # An exception should be raised if the backend isn't specified but + # options are given which are not an instance of any of the known + # agents' option classes. + rpc_backend_options = FooBackendOptions(self.init_method) + + with self.assertRaisesRegex(TypeError, "Could not infer backend for options"): + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + # Do _not_ pass backend. + rpc_backend_options=rpc_backend_options, + ) + + @dist_init + def test_owner_rref_backward(self): + dst = worker_name((self.rank + 1) % self.world_size) + t1 = torch.rand(10, 10, requires_grad=True) + rref = rpc.RRef(t1.sum() + t1.sum()) + rref.backward() + expected_grad = torch.ones_like(t1) * 2 + self.assertEqual(expected_grad, t1.grad) + + with dist_autograd.context() as context_id: + t2 = rpc.rpc_sync(dst, torch.add, args=(t1, t1)) + rref = rpc.RRef(t2.sum()) + rref.backward(context_id) + self.assertEqual(expected_grad, dist_autograd.get_gradients(context_id)[t1]) + + # Double backward. + with dist_autograd.context() as context_id: + t2 = rpc.rpc_sync(dst, torch.add, args=(t1, t1)) + rref = rpc.RRef(t2.sum()) + rref.backward(context_id, retain_graph=True) + rref.backward(context_id) + self.assertEqual(expected_grad * 2, dist_autograd.get_gradients(context_id)[t1]) + + # Test errors. + with self.assertRaisesRegex(RuntimeError, "tensors does not require grad and does not have a grad_fn"): + rpc.RRef(torch.rand(10)).backward() + + with self.assertRaisesRegex(RuntimeError, "grad can be implicitly created only for scalar outputs"): + rpc.RRef(torch.rand(10, requires_grad=True)).backward() + + with self.assertRaisesRegex(RuntimeError, "Could not find autograd context with id: 100"): + rpc.RRef(torch.rand(10, requires_grad=True).sum()).backward(100) + + with self.assertRaisesRegex(RuntimeError, "RRef should contain a tensor for .backward()"): + rpc.RRef("foo").backward() + + @staticmethod + def _sum(x): + return x.sum() + + @staticmethod + def _identity(x): + return x + + @dist_init + def test_user_rref_backward(self): + dst = worker_name((self.rank + 1) % self.world_size) + t = torch.rand(10, requires_grad=True) + with dist_autograd.context() as context_id: + rref = rpc.remote(dst, RpcTest._sum, args=(t,)) + rref.backward(context_id, retain_graph=True) + rref.backward(context_id) + self.assertEqual(torch.ones_like(t) * 2, dist_autograd.get_gradients(context_id)[t]) + + with dist_autograd.context() as context_id: + rref = rpc.remote(dst, RpcTest._identity, args=("foo",)) + with self.assertRaisesRegex(RuntimeError, "RRef should contain a tensor for .backward()"): + rref.backward(context_id) + + with self.assertRaisesRegex(RuntimeError, "User RRefs require 'dist_autograd_ctx_id' to be specified"): + rref.backward() + + @dist_init(setup_rpc=False) + def test_shutdown_errors(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + if self.rank != 0: + og_func = rpc.api._broadcast_to_followers + og_rref_func = rpc.api._delete_all_user_and_unforked_owner_rrefs + + # Monkey-patch _broadcast_to_followers to fail, which would ensure + # _all_gather on leader raises an exception. + def raise_error(sequence_id, objects_map): + og_func(sequence_id, objects_map) + raise RuntimeError('simulation') + + # Monkey-patch _delete_all_user_and_unforked_owner_rrefs to fail, + # which would ensure barrier is not called on followers. + def rref_error(): + raise RuntimeError('simulation rref') + + try: + rpc.api._broadcast_to_followers = raise_error + rpc.api._delete_all_user_and_unforked_owner_rrefs = rref_error + with self.assertRaisesRegex(RuntimeError, 'simulation rref'): + rpc.shutdown() + finally: + rpc.api._broadcast_to_followers = og_func + rpc.api._delete_all_user_and_unforked_owner_rrefs = og_rref_func + else: + with self.assertRaisesRegex(RuntimeError, 'timed out in _all_gather'): + rpc.shutdown() + + dist.barrier() + + @dist_init + def test_my_parameter_server(self): + self._my_parameter_server(False) + + +class CudaRpcTest(RpcAgentTestFixture): + + @skip_if_lt_x_gpu(2) + @dist_init + def test_profiler_remote_cuda(self): + if self.rank != 1: + return + + dst_cuda_0 = (self.rank + 1) % self.world_size + dst_cuda_1 = (self.rank + 2) % self.world_size + dst_worker_cuda_0 = worker_name(dst_cuda_0) + dst_worker_cuda_1 = worker_name(dst_cuda_1) + + with _profile(use_cuda=True) as p: + fut1 = rpc.rpc_async(dst_worker_cuda_0, udf_with_torch_ops, args=(0, )) + fut2 = rpc.rpc_async(dst_worker_cuda_1, udf_with_torch_ops, args=(1, )) + fut1.wait() + fut2.wait() + + def get_name(event): + return event.name[event.name.find(REMOTE_OP_STR) + len(REMOTE_OP_STR):] + + function_events = p.function_events + for event in function_events: + if event.is_async: + self.assertEqual(0, event.device_time_total) + self.assertEqual([], event.kernels) + self.assertEqual(0, event.device_time) + else: + if event.node_id == 1: + continue + self.assertTrue(event.node_id in [dst_cuda_0, dst_cuda_1]) + if get_name(event) in EXPECTED_REMOTE_EVENTS: + self.assertGreater(event.device_time_total, 0) + self.assertEqual(1, len(event.kernels)) + kernel = event.kernels[0] + if event.node_id == dst_cuda_0: + self.assertEqual(kernel.device, 0) + if event.node_id == dst_cuda_1: + self.assertEqual(kernel.device, 1) + self.assertGreater(event.device_time, 0) + + # Validate that EXPECTED_REMOTE_EVENTS is a subset of remotely profiled + # events. + remote_events = [event for event in function_events if event.is_remote] + remote_event_names = [get_name(event) for event in remote_events if get_name(event) in EXPECTED_REMOTE_EVENTS] + self.assertEqual(set(remote_event_names), set(EXPECTED_REMOTE_EVENTS)) + + +class TensorPipeAgentRpcTest(RpcAgentTestFixture, RpcTestCommon): + + def test_mismatched_type_for_options(self): + # An exception should be raised if the options are not an instance of + # TensorPipeRpcBackendOptions. + rpc_backend_options = FooBackendOptions(self.init_method) + + with self.assertRaisesRegex( + TypeError, "`rpc_backend_options` must be a `TensorPipeRpcBackendOptions`" + ): + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + backend=rpc.BackendType.TENSORPIPE, + rpc_backend_options=rpc_backend_options, + ) + + def test_infer_backend_from_options(self): + rpc_backend_options = rpc.TensorPipeRpcBackendOptions( + init_method=self.init_method, + _transports=tp_transports() + ) + + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + # Do _not_ pass backend. + rpc_backend_options=rpc_backend_options, + ) + + self.assertIsInstance(rpc.api._get_current_rpc_agent(), rpc.TensorPipeAgent) + + # FIXME Merge this test with the corresponding one in RpcTest. + @dist_init(setup_rpc=False) + def test_set_and_get_num_worker_threads(self): + NUM_THREADS = 27 + rpc_backend_options = rpc.TensorPipeRpcBackendOptions( + init_method=self.rpc_backend_options.init_method, + num_worker_threads=NUM_THREADS, + _transports=tp_transports(), + ) + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + + info = rpc.api._get_current_rpc_agent().get_debug_info() + self.assertEqual(int(info["agent.thread_pool_size"]), NUM_THREADS) + rpc.shutdown() + + # FIXME Merge this test with the corresponding one in RpcTest. + @dist_init(setup_rpc=False) + def test_tensorpipe_set_default_timeout(self): + # Set a high timeout since it doesn't affect test runtime and ensures + # the test doesn't erroneously timeout due to slow machines. + timeout = 100 + rpc_backend_options = rpc.TensorPipeRpcBackendOptions( + init_method=self.rpc_backend_options.init_method, + num_worker_threads=self.rpc_backend_options.num_worker_threads, + rpc_timeout=timeout, + _transports=tp_transports(), + ) + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + + default_timeout = rpc.get_rpc_timeout() + self.assertEqual(default_timeout, timeout) + rpc.shutdown() + + # FIXME Merge this test with the corresponding one in RpcTest. + @dist_init(setup_rpc=False) + def test_tensorpipe_options_throw_on_timedelta_timeout(self): + from datetime import timedelta + + timeout = timedelta() + # Ensure that constructing TensorPipeRpcBackendOptions with timedelta fails + with self.assertRaisesRegex(TypeError, "incompatible constructor arguments"): + rpc.TensorPipeRpcBackendOptions( + init_method=self.rpc_backend_options.init_method, + num_worker_threads=self.rpc_backend_options.num_worker_threads, + rpc_timeout=timeout, + ) + + @dist_init + def _test_rref_get_type_timeout(self, blocking): + # Test where we try to get the type of a RRef from an owner, but RRef + # creation is slower than timeout passed into _get_type. + dst_rank = (self.rank + 1) % self.world_size + dst = worker_name(dst_rank) + slow_rref = rpc.remote(dst, MyClass, args=(torch.ones(2, 2), True)) + timeout = 0.5 + expected_err = self.get_timeout_error_regex() + # Blocking: blocks on inline call + if blocking: + with self.assertRaisesRegex(RuntimeError, expected_err): + slow_rref._get_type(timeout=timeout, blocking=blocking) + # Non-blocking: blocks on wait + else: + fut = slow_rref._get_type(timeout=timeout, blocking=blocking) + with self.assertRaisesRegex(RuntimeError, expected_err): + fut.wait() + + # FIXME We wait until the remote completed creating the OwnerRRef + # because there's currently a race if we shut down RPC before that. + slow_rref.to_here() + + def test_rref_get_type_timeout_blocking(self): + self._test_rref_get_type_timeout(blocking=True) + + def test_rref_get_type_timeout_non_blocking(self): + self._test_rref_get_type_timeout(blocking=False) + + @dist_init + def test_op_with_invalid_args(self): + dst = worker_name((self.rank + 1) % self.world_size) + with self.assertRaisesRegex( + RuntimeError, "Overloaded torch operator invoked from Python failed to match any schema" + ): + rpc.rpc_sync(dst, torch.add, args=()) + + def _test_rref_proxy_timeout(self, rref_proxy_api): + dst_rank = (self.rank + 1) % self.world_size + dst = worker_name(dst_rank) + rref = rpc.remote(dst, MyClass, args=(torch.ones(2, 2), )) + # Ensure RRef is created on remote node. + rref.to_here() + rref_api = getattr(rref, rref_proxy_api) + self.assertTrue(rref_api is not None, f"Failed to get RRef proxy api: {rref_proxy_api}") + expected_error = self.get_timeout_error_regex() + timeout = 2 + with self.assertRaisesRegex(RuntimeError, expected_error): + result = rref_api(timeout=timeout).my_slow_method(torch.ones(2, 2)) + if rref_api == rref.rpc_async: + result.wait() + elif rref_api == rref.remote: + result._get_future().wait() + + # Case where rpc.remote() is stuck and exceeds timeout + slow_rref = rpc.remote(dst, MyClass, args=(torch.ones(2, 2), True)) + timeout = 0.01 + rref_api = getattr(slow_rref, rref_proxy_api) + # Note that even when we call rref.rpc_async() in this case, we + # time out in future creation, not waiting for future. This is because + # rref proxy function calls rref._get_type before returning future, + # which blocks on the RRef being created on owner node, until the + # specified timeout. + with self.assertRaisesRegex(RuntimeError, expected_error): + result = rref_api(timeout=timeout).my_instance_method(torch.ones(2, 2)) + # rpc_async returns immediately and surface a timeout through wait() + if rref_api == slow_rref.rpc_async: + result.wait() + + # FIXME We wait until the remote completed creating the OwnerRRef + # because there's currently a race if we shut down RPC before that. + slow_rref.to_here() + + @dist_init + def test_rref_proxy_timeout(self): + for rpc_api in ["rpc_sync", "rpc_async", "remote"]: + self._test_rref_proxy_timeout(rpc_api) + + @dist_init + def test_send_to_rank_sparse(self): + dst_rank = (self.rank + 1) % self.world_size + + # Test sparse tensor + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + x = build_sparse_tensor() + y = build_sparse_tensor() + expected_tensor = (x + y) + ret = self._run_func_in_mode(dst_rank, torch.add, exec_mode, args=(x, y)) + self.assertEqual(expected_tensor, ret) + + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + x = build_sparse_tensor(coalesce=True) + y = build_sparse_tensor(coalesce=True) + expected_tensor = (x + y) + ret = self._run_func_in_mode(dst_rank, torch.add, exec_mode, args=(x, y)) + self.assertEqual(expected_tensor, ret) + + @dist_init + def test_self_py_udf_remote_sparse(self): + self._self_py_udf_remote( + rpc.get_worker_info(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() + ) + + @dist_init + def test_self_remote_rref_as_rpc_arg_sparse(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._self_remote_rref_as_rpc_arg( + dst, + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() + ) + + @dist_init + def test_self_remote_rref_as_self_rpc_arg_sparse(self): + self._self_remote_rref_as_rpc_arg( + rpc.get_worker_info(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() + ) + + @dist_init + def test_self_remote_rref_as_remote_arg_sparse(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._self_remote_rref_as_remote_arg( + dst, + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() + ) + + @dist_init + def test_self_remote_rref_as_self_remote_arg_sparse(self): + self._self_remote_rref_as_remote_arg( + rpc.get_worker_info(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() + ) + + def test_world_size_one_sparse(self): + self._world_size_one( + build_sparse_tensor(), + build_sparse_tensor() + ) + + @dist_init + def test_multi_rpc_sparse(self): + self._multi_rpc(True) + + def test_wait_all_workers_sparse(self): + self._wait_all_workers(heavy_rpc_sparse, build_sparse_tensor()) + + def test_wait_all_workers_twice_sparse(self): + self._wait_all_workers_twice(heavy_rpc_sparse, build_sparse_tensor()) + + @dist_init + def test_py_sparse_tensors_in_container(self): + n = self.rank + 1 + dst_rank = n % self.world_size + a = [build_sparse_tensor(), build_sparse_tensor()] + ret = rpc.rpc_sync( + worker_name(dst_rank), my_container_sum, args=(a,) + ) + self.assertEqual(ret, my_container_sum(a)) + + @dist_init + def test_nested_rpc_sparse(self): + self._nested_rpc(nested_rpc_sparse, build_sparse_tensor() * 2) + + @dist_init + def test_stress_heavy_rpc_sparse(self): + self._stress_test_rpc(heavy_rpc_sparse, repeat=20, args=(build_sparse_tensor(),)) + + @dist_init + def test_builtin_remote_ret_sparse(self): + self._builtin_remote_ret( + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() * 2 + ) + + @dist_init + def test_builtin_remote_self_sparse(self): + self._builtin_remote_self( + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() * 2 + ) + + @dist_init + def test_multi_builtin_remote_ret_sparse(self): + self._test_multi_remote_call( + torch.add, True, + args_fn=RpcTest._multi_args_fn + ) + + @dist_init + def test_multi_py_udf_remote_sparse(self): + self._test_multi_remote_call( + my_function, + True, + kwargs_fn=RpcTest._multi_kwargs_fn + ) + + @dist_init + def test_py_rref_args_sparse(self): + self._py_rref_args( + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() * 4 + ) + + @dist_init + def test_py_rref_args_user_share_sparse(self): + self._py_rref_args_user_share( + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() * 6 + ) + + @dist_init + def test_py_rpc_rref_args_sparse(self): + self._py_rpc_rref_args( + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() * 6 + ) + + @dist_init + def test_nested_remote_sparse(self): + self._nested_remote( + nested_remote_sparse, + build_sparse_tensor() + build_sparse_tensor() + ) + + @dist_init + def test_nested_rref_sparse(self): + self._nested_rref( + nested_rref_sparse, + build_sparse_tensor() * 2, + build_sparse_tensor() * 2 + ) + + @dist_init + def test_nested_rref_stress_sparse(self): + self._nested_rref_stress( + nested_rref_sparse, + build_sparse_tensor() * 2, + build_sparse_tensor() * 2 + ) + + @dist_init + def test_my_parameter_server_sparse(self): + self._my_parameter_server(True) + + # Test init_rpc without world_size argument + @dist_init(setup_rpc=False) + def test_dynamic_rpc_init_rpc(self): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + rpc.shutdown() + + # Dynamic RPC new ranks communicate with existing ranks + @dist_init(setup_rpc=False) + def test_dynamic_rpc_new_rank_can_communicated_with_existing_rank(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + if self.rank == 0: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + + # Rank 0 will be initialized with RPC after this barrier + dist.barrier() + + if self.rank != 0: + # Newly joined ranks will be able to communicate with rank 0, since that was created first + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + result = rpc.rpc_sync(worker_name(0), torch.add, args=(torch.tensor(1), torch.tensor(1))) + self.assertEqual(torch.add(torch.tensor(1), torch.tensor(1)), result) + + # Barrier to ensure that all rpc_sync calls are finished + dist.barrier() + rpc.shutdown() + + # Dynamic RPC existing ranks can communicate with new ranks + @dist_init(setup_rpc=False) + def test_dynamic_rpc_existing_rank_can_communicate_with_new_rank(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + if self.rank == 0: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + + # Rank 0 will be initialized with RPC after this barrier + dist.barrier() + + # Rest of ranks join after barrier + if self.rank != 0: + # Newly joined ranks will be able to communicate with rank 0, since that was created first + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + + dist.barrier() + if self.rank == 0: + for i in range(1, self.world_size): + result = rpc.rpc_sync(worker_name(i), torch.add, args=(torch.tensor(1), torch.tensor(1))) + self.assertEqual(torch.add(torch.tensor(1), torch.tensor(1)), result) + + # Barrier to ensure that all rpc_sync calls are finished + dist.barrier() + rpc.shutdown() + + # Dynamic RPC existing ranks can communicate with new ranks using CUDA rpc + @skip_if_lt_x_gpu(2) + @dist_init(setup_rpc=False) + def test_dynamic_rpc_existing_rank_can_communicate_with_new_rank_cuda(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + if self.rank == 0: + options = self.rpc_backend_options + for i in range(1, self.world_size): + dst = worker_name(i) + options.set_device_map(dst, {1: 0}) + options.set_device_map(dst, {0: 1}) + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=options, + ) + + # Rank 0 will be initialized with RPC after this barrier + dist.barrier() + + # Rest of ranks join after barrier + if self.rank != 0: + # Newly joined ranks will be able to communicate with rank 0, since that was created first + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + + # TODO: Cuda RPC is failing due to: + # terminate called after throwing an instance of 'c10::Error' + # what(): 0 <= device && static_cast(device) < device_allocator.size() + # INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":1937, + # please report a bug to PyTorch. Allocator not initialized for device 1: did you call init? + # dist.barrier() + # if self.rank == 0: + # for i in range(1, self.world_size): + # x = torch.ones(2) + # result_on_device_0 = rpc.rpc_sync(worker_name(i), torch.add, args=(x.to(0), 1)) + # result_on_device_1 = rpc.rpc_sync(worker_name(i), torch.add, args=(x.to(1), 1)) + # self.assertEqual(torch.add(torch.ones(2), 1), result_on_device_0) + # self.assertEqual(torch.device('cuda:0'), result_on_device_0.device) + # self.assertEqual(torch.add(torch.ones(2), 1), result_on_device_1) + # self.assertEqual(torch.device('cuda:1'), result_on_device_1.device) + + # Barrier to ensure that all rpc_sync calls are finished + dist.barrier() + rpc.shutdown() + + @dist_init(setup_rpc=False) + def test_dynamic_rpc_init_rpc_without_rank(self): + # default initialization uses file init + with self.assertRaisesRegex(ValueError, "rank parameter missing"): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rpc_backend_options=self.rpc_backend_options, + ) + + # env init + with self.assertRaisesRegex(ValueError, "environment variable RANK expected"): + rpc_backend_options = rpc.TensorPipeRpcBackendOptions(init_method="env://") + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rpc_backend_options=rpc_backend_options, + ) + + # tcp init + with self.assertRaisesRegex(ValueError, "rank parameter missing"): + rpc_backend_options = rpc.TensorPipeRpcBackendOptions(init_method="tcp://127.0.0.1:23456") + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rpc_backend_options=rpc_backend_options, + ) + + @dist_init(setup_rpc=False) + def test_dynamic_and_static_init_rpc_together(self): + # Initialize a static rpc group with size = self.world_size - 1 + dist.init_process_group( + backend='gloo', + init_method=self.file_init_method, + rank=self.rank, + world_size=self.world_size) + + world_size_minus_one = self.world_size - 1 + if self.rank < world_size_minus_one: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=world_size_minus_one, + rpc_backend_options=self.rpc_backend_options, + ) + + dist.barrier() + + # Attempt to add an additional dynamic group member + if self.rank == world_size_minus_one: + # Expect error message to be thrown + with self.assertRaisesRegex(RuntimeError, "RPC group mixes statically and dynamically\ + initialized members which is not supported."): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + +class TensorPipeAgentCudaRpcTest(RpcAgentTestFixture, RpcTestCommon): + + def _test_device_maps(self, options, errMsg): + with self.assertRaisesRegex(ValueError, errMsg): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + self.assertFalse(rpc.api._is_current_rpc_agent_set()) + + @skip_if_lt_x_gpu(2) + def test_device_maps_wrong_worker_name(self): + options = self.rpc_backend_options + options.set_device_map("none_exist", {0: 1}) + + self._test_device_maps( + options, + errMsg="Node worker0 has invalid target node names in its device maps" + ) + + @skip_if_lt_x_gpu(1) + def test_device_maps_invalid_max_local_device(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {torch.cuda.device_count(): 0}) + + self._test_device_maps( + options, + errMsg="Node worker0 has source devices with invalid indices in its device map for worker1" + ) + + @skip_if_lt_x_gpu(1) + def test_device_maps_invalid_max_remote_device(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {0: torch.cuda.device_count()}) + + self._test_device_maps( + options, + errMsg="Node worker0 has target devices with invalid indices in its device map for worker1" + ) + + @skip_if_lt_x_gpu(2) + def test_device_maps_many_to_one(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {1: 0}) + options.set_device_map(dst, {0: 0}) + + self._test_device_maps( + options, + errMsg="Node worker0 has duplicated target devices in its device map for worker1" + ) + + @skip_if_lt_x_gpu(2) + def test_device_maps_one_to_many(self): + if self.rank == 0: + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {0: 1}) + with self.assertRaisesRegex( + ValueError, "`set_device_map` only supports 1-to-1 mapping" + ): + options.set_device_map(dst, {0: 0}) + + @skip_if_lt_x_gpu(1) + def test_device_maps_invalid_min_device(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + with self.assertRaisesRegex( + RuntimeError, "Device index must not be negative" + ): + options.set_device_map(dst, {-1: 0}) + + with self.assertRaisesRegex( + RuntimeError, "Device index must not be negative" + ): + options.set_device_map(dst, {0: -1}) + + @staticmethod + def _gpu_add(x, y): + if all([x.is_cuda, x.device.index == 1, y.is_cuda, y.device.index == 1]): + return (x + y).to(0) + else: + raise ValueError("Wrong device affinity") + + @skip_if_lt_x_gpu(2) + def test_device_maps_gpu(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {0: 1, 1: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + ret = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._gpu_add, + args=(torch.zeros(2).to(0), torch.ones(2).to(0)) + ) + self.assertEqual(ret.device, torch.device(1)) + self.assertEqual(ret, (torch.zeros(2) + torch.ones(2)).to(1)) + rpc.shutdown() + + @staticmethod + def _gpu_add_given_devices(x, y, x_to, y_to, z_to): + x_device = "cpu" if x.device.type == "cpu" else x.device.index + y_device = "cpu" if y.device.type == "cpu" else y.device.index + if x_device == x_to and y_device == y_to: + return x.to(z_to) + y.to(z_to) + else: + raise ValueError("Wrong device affinity") + + def _test_device_maps_gpu(self, x_from, y_from, z_to, device_map, dst=None, fn=None): + fn = TensorPipeAgentCudaRpcTest._gpu_add_given_devices if fn is None else fn + x_to = device_map[x_from] + y_to = device_map[y_from] + + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) if dst is None else dst + options.set_device_map(dst, device_map) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + x = torch.zeros(2).to(x_from) + y = torch.ones(2).to(y_from) + + ret = rpc.rpc_sync(dst, fn, args=(x, y, x_to, y_to, z_to)) + + reverse_device_map = {device_map[k] : k for k in device_map} + z_from = reverse_device_map[z_to] + + ret_device = "cpu" if ret.device.type == "cpu" else ret.device.index + self.assertEqual(ret_device, z_from) + self.assertEqual(ret, torch.ones(2).to(z_from)) + + rpc.shutdown() + + def test_device_map_cpu(self): + self._test_device_maps_gpu( + x_from="cpu", + y_from="cpu", + z_to="cpu", + device_map={"cpu" : "cpu"}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(1) + def test_device_map_cpu_to_gpu_default(self): + self._test_device_maps_gpu( + x_from="cpu", + y_from="cpu", + z_to=0, + device_map={"cpu" : 0}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_cpu_to_gpu_non_default(self): + self._test_device_maps_gpu( + x_from="cpu", + y_from="cpu", + z_to=1, + device_map={"cpu" : 1}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(1) + def test_device_map_gpu_to_cpu_default(self): + self._test_device_maps_gpu( + x_from=0, + y_from=0, + z_to="cpu", + device_map={0 : "cpu"}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_to_cpu_non_default(self): + self._test_device_maps_gpu( + x_from=1, + y_from=1, + z_to="cpu", + device_map={1 : "cpu"}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_default(self): + self._test_device_maps_gpu( + x_from=0, + y_from=0, + z_to=0, + device_map={0 : 0} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_non_default(self): + self._test_device_maps_gpu( + x_from=1, + y_from=1, + z_to=1, + device_map={1 : 1} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_default_to_non_default(self): + self._test_device_maps_gpu( + x_from=0, + y_from=0, + z_to=1, + device_map={0 : 1} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_non_default_to_default(self): + self._test_device_maps_gpu( + x_from=1, + y_from=1, + z_to=0, + device_map={1 : 0} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_1(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=0, + device_map={0 : 0, 1 : 1} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_2(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=1, + device_map={0 : 0, 1 : 1} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_3(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=0, + device_map={0 : 0, 1 : 1} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_4(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=1, + device_map={0 : 0, 1 : 1} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_5(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=0, + device_map={0 : 1, 1 : 0} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_6(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=1, + device_map={0 : 1, 1 : 0} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_7(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=0, + device_map={0 : 1, 1 : 0} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_8(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=1, + device_map={0 : 1, 1 : 0} + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_1(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=0, + device_map={0 : 0, 1 : 1}, + dst=worker_name(self.rank) + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_2(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=1, + device_map={0 : 0, 1 : 1}, + dst=worker_name(self.rank) + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_3(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=0, + device_map={0 : 0, 1 : 1}, + dst=worker_name(self.rank) + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_4(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=1, + device_map={0 : 0, 1 : 1}, + dst=worker_name(self.rank) + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_5(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=0, + device_map={0 : 1, 1 : 0}, + dst=worker_name(self.rank) + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_6(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=1, + device_map={0 : 1, 1 : 0}, + dst=worker_name(self.rank) + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_7(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=0, + device_map={0 : 1, 1 : 0}, + dst=worker_name(self.rank) + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_8(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=1, + device_map={0 : 1, 1 : 0}, + dst=worker_name(self.rank) + ) + + @staticmethod + def _gpu_add_multi_gpu(x, y): + if all([x.is_cuda, x.device.index == 1, y.is_cuda, y.device.index == 0]): + return x.to(0) + y, x - y.to(1) + else: + raise ValueError("Wrong device affinity") + + def _test_device_maps_multi_gpu(self, dst): + options = self.rpc_backend_options + options.set_device_map(dst, {0: 1}) + options.set_device_map(dst, {1: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + x = torch.zeros(2).to(0) + y = torch.ones(2).to(1) + rets = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._gpu_add_multi_gpu, + args=(x, y) + ) + + self.assertEqual(rets[0].device, torch.device(1)) + self.assertEqual(rets[1].device, torch.device(0)) + self.assertEqual(rets[0], (torch.zeros(2) + torch.ones(2)).to(1)) + self.assertEqual(rets[1], (torch.zeros(2) - torch.ones(2)).to(0)) + rpc.shutdown() + + @skip_if_lt_x_gpu(2) + def test_device_maps_multi_gpu(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._test_device_maps_multi_gpu(dst) + + @skip_if_lt_x_gpu(2) + def test_device_maps_multi_gpu_self(self): + dst = worker_name(self.rank) + self._test_device_maps_multi_gpu(dst) + + @staticmethod + def _gpu_add_return_to_gpu(x, y): + if x.device.type == 'cpu' and y.device.type == 'cpu': + return (x + y).to(0), (x - y).to(1), (x * y).to(2), (x / y).to(3) + else: + raise ValueError("Wrong device affinity") + + @skip_if_lt_x_gpu(2) + def test_device_maps_in_options(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc.TensorPipeRpcBackendOptions( + init_method=options.init_method, + num_worker_threads=options.num_worker_threads, + device_maps={dst: {0: 1, 1: 0}}, + _transports=tp_transports() + ) + ) + + rets = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._gpu_add_multi_gpu, + args=(torch.zeros(2).to(0), torch.ones(2).to(1)) + ) + self.assertEqual(rets[0].device, torch.device(1)) + self.assertEqual(rets[1].device, torch.device(0)) + self.assertEqual(rets[0], (torch.zeros(2) + torch.ones(2)).to(1)) + self.assertEqual(rets[1], (torch.zeros(2) - torch.ones(2)).to(0)) + rpc.shutdown() + + def _test_device_maps_return_to_gpu(self, dst): + options = self.rpc_backend_options + + options.set_device_map(dst, {0: 1}) + options.set_device_map(dst, {1: 2}) + options.set_device_map(dst, {2: 3}) + options.set_device_map(dst, {3: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rets = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._gpu_add_return_to_gpu, + args=(torch.zeros(2), torch.ones(2)) + ) + for i in range(len(rets)): + self.assertEqual(rets[i].device, torch.device((3 + i) % 4)) + self.assertEqual(rets[0], (torch.zeros(2) + torch.ones(2)).to(3)) + self.assertEqual(rets[1], (torch.zeros(2) - torch.ones(2)).to(0)) + self.assertEqual(rets[2], (torch.zeros(2) * torch.ones(2)).to(1)) + self.assertEqual(rets[3], (torch.zeros(2) / torch.ones(2)).to(2)) + rpc.shutdown() + + @skip_if_lt_x_gpu(4) + def test_device_maps_return_to_gpu(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._test_device_maps_return_to_gpu(dst) + + @skip_if_lt_x_gpu(4) + def test_device_maps_return_to_gpu_self(self): + dst = worker_name(self.rank) + self._test_device_maps_return_to_gpu(dst) + + @staticmethod + def _add_to_gpu(x, y): + return (x + y).to(0) + + def _test_device_maps_missing_config(self, mode): + dst = worker_name((self.rank + 1) % self.world_size) + errMsg = ( + "TensorPipe RPC backend only supports CPU tensors by default.*" + "`set_device_map` on `TensorPipeRpcBackendOptions`" + ) + + with self.assertRaisesRegex(RuntimeError, errMsg): + if mode == RPCExecMode.SYNC: + rpc.rpc_sync(dst, torch.add, args=(torch.zeros(2).to(0), 1)) + elif mode == RPCExecMode.REMOTE: + rpc.remote(dst, torch.add, args=(torch.zeros(2).to(0), 1)).to_here() + else: + raise ValueError(f"unexpected mode {mode}") + + # make sure RPC is still functioning + ret = rpc.rpc_sync(dst, torch.add, args=(torch.ones(2), 1)) + self.assertEqual(ret, torch.ones(2) + 1) + + def _test_device_maps_missing_config_response(self, mode): + dst = worker_name((self.rank + 1) % self.world_size) + errMsg = "Response device mapping is not available" + + with self.assertRaisesRegex(RuntimeError, errMsg): + if mode == RPCExecMode.SYNC: + rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._add_to_gpu, + args=(torch.zeros(2), 1) + ) + elif mode == RPCExecMode.REMOTE: + rpc.remote( + dst, + TensorPipeAgentCudaRpcTest._add_to_gpu, + args=(torch.zeros(2), 1) + ).to_here() + else: + raise ValueError(f"unexpected mode {mode}") + + # make sure RPC is still functioning + ret = rpc.rpc_sync(dst, torch.add, args=(torch.ones(2), 1)) + self.assertEqual(ret, torch.ones(2) + 1) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config(self): + self._test_device_maps_missing_config(RPCExecMode.SYNC) + + @skip_if_lt_x_gpu(1) + def test_device_maps_missing_config_not_timeout(self): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options + ) + + timeout = rpc.get_rpc_timeout() + + tik = time.time() + self._test_device_maps_missing_config(RPCExecMode.SYNC) + rpc.shutdown() + tok = time.time() + + self.assertTrue(tok - tik < timeout) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_loop(self): + for _ in range(self.rpc_backend_options.num_worker_threads + 5): + self._test_device_maps_missing_config(RPCExecMode.SYNC) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_response(self): + self._test_device_maps_missing_config_response(RPCExecMode.SYNC) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_response_loop(self): + for _ in range(self.rpc_backend_options.num_worker_threads + 5): + self._test_device_maps_missing_config_response(RPCExecMode.SYNC) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_remote(self): + self._test_device_maps_missing_config(RPCExecMode.REMOTE) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_remote_response(self): + self._test_device_maps_missing_config_response(RPCExecMode.REMOTE) + + @skip_if_lt_x_gpu(2) + def test_device_maps_remote(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {1: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rref = rpc.remote( + dst, + TensorPipeAgentCudaRpcTest._add_to_gpu, + args=(torch.zeros(2), 1) + ) + + self.assertEqual(rref.to_here().device.index, 1) + self.assertEqual(rref.to_here(), torch.ones(2).to(1)) + + rpc.shutdown() + + @staticmethod + def _slow_add_on_user_stream(x, y): + s0 = torch.cuda.current_stream(x.device) + s1 = torch.cuda.Stream(device=x.device) + s1.wait_stream(s0) + x.record_stream(s1) + y.record_stream(s1) + with torch.cuda.stream(s1): + torch.cuda._sleep(10 * FIFTY_MIL_CYCLES) + z = x + y + s0.wait_stream(s1) + z.record_stream(s0) + return z + + def _test_custom_stream(self, fn, device_map): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, device_map) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + fn(dst) + + rpc.shutdown() + + def _test_stream_sync(self, dst): + x = torch.ones(2, 2).to(0) + ret = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._slow_add_on_user_stream, + args=(x, x) + ) + self.assertEqual(ret, 2 * x) + + @skip_if_lt_x_gpu(2) + def test_custom_stream(self): + self._test_custom_stream(self._test_stream_sync, {"cuda:0": "cuda:1"}) + + def _test_stream_multi_async(self, dst): + futs = [] + for i in range(20): + x = torch.ones(2, 2).to(0) * i + futs.append( + rpc.rpc_async( + dst, + TensorPipeAgentCudaRpcTest._slow_add_on_user_stream, + args=(x, x) + ) + ) + + for i in range(20): + self.assertEqual(futs[i].wait(), 2 * torch.ones(2, 2).to(0) * i) + + @skip_if_lt_x_gpu(2) + def test_custom_stream_multi(self): + self._test_custom_stream( + self._test_stream_multi_async, + {"cuda:0": "cuda:1"} + ) + + @staticmethod + def _nested_slow_add_on_user_stream(dst, x, y, z): + ret = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._slow_add_on_user_stream, + args=(x, y) + ) + + return TensorPipeAgentCudaRpcTest._slow_add_on_user_stream(ret, z) + + def _test_stream_nested_sync(self, dst): + x = torch.ones(2, 2).to(0) + y = torch.ones(2, 2).to(0) * 2 + z = torch.ones(2, 2).to(0) * 3 + nested_dst = worker_name((self.rank + 2) % self.world_size) + ret = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._nested_slow_add_on_user_stream, + args=(nested_dst, x, y, z) + ) + self.assertEqual(ret, 6 * x) + + @skip_if_lt_x_gpu(2) + def test_custom_stream_nested(self): + self._test_custom_stream( + self._test_stream_nested_sync, + {"cuda:0": "cuda:1", "cuda:1": "cuda:0"} + ) + + def _test_stream_nested_multi_async(self, dst): + if self.rank == 0: + futs = [] + n = 5 + xs, ys, zs = [], [], [] + for i in range(n): + x = torch.ones(2, 2).to(0) * (i - 1) + y = torch.ones(2, 2).to(0) * i + z = torch.ones(2, 2).to(0) * (i + 1) + xs.append(x) + ys.append(y) + zs.append(z) + nested_dst = worker_name((self.rank + 2) % self.world_size) + futs.append( + rpc.rpc_async( + dst, + TensorPipeAgentCudaRpcTest._nested_slow_add_on_user_stream, + args=(nested_dst, x, y, z) + ) + ) + + for i in range(n): + self.assertEqual(futs[i].wait(), xs[i] + ys[i] + zs[i]) + + @skip_if_lt_x_gpu(2) + def test_custom_stream_nested_multi(self): + self._test_custom_stream( + self._test_stream_nested_multi_async, + {"cuda:0": "cuda:1", "cuda:1": "cuda:0"} + ) + + @staticmethod + def _gpu_add_wrong_gpus(x, y): + if x.is_cuda and y.is_cuda: + return x.cpu() + y.cuda() + else: + raise ValueError("Wrong device affinity") + + @skip_if_lt_x_gpu(1) + def test_device_mismatch(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {0: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + x = torch.zeros(2).to(0) + y = torch.ones(2).to(0) + + with self.assertRaisesRegex( + RuntimeError, + "Expected all tensors to be on the same device, but found at least two devices" + ): + rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._gpu_add_wrong_gpus, + args=(x, y) + ) + + rpc.shutdown() + + def _test_rref_synchronization(self, local_device, remote_device): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {local_device : remote_device}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + if self.rank == 1: + # This test compares rref.rpc_sync().forward(x) vs rref.remote().forward(x).to_here() + # If to_here() is properly synchronized with forward(x) the results must be identical + # This test needs multiple iterations and significant batch size to simulate real + # training of a CNN of MNIST-like data. + # see https://github.com/pytorch/pytorch/issues/54771 + rref = rpc.remote(dst, MyConvNetForMNIST, args=(remote_device,)) + for _ in range(10): + x = torch.randn(200, 1, 28, 28).to(local_device) + actual = rref.remote().forward(x).to_here() + expected = rref.rpc_sync().forward(x) + self.assertEqual(actual, expected) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_rref_to_here_synchronization1(self): + self._test_rref_synchronization("cuda:0", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_to_here_synchronization2(self): + self._test_rref_synchronization("cuda:1", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_to_here_synchronization3(self): + self._test_rref_synchronization("cuda:1", "cuda:1") + + @skip_if_lt_x_gpu(2) + def test_rref_to_here_synchronization4(self): + self._test_rref_synchronization("cuda:0", "cuda:1") + + def _test_rref_as_arg_synchronization( + self, + local_device, + remote_device, + devicesOptions=None + ): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {local_device: remote_device}) + + input_src = worker_name((self.rank - 1 + self.world_size) % self.world_size) + options.set_device_map(input_src, {remote_device: local_device}) + + if devicesOptions is not None: + options.set_devices(devicesOptions[self.rank]) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + if self.rank == 1: + # This test compares rref.rpc_sync().forward(x) vs rref.remote().forward(x).to_here() + # If to_here() is properly synchronized with forward(x) the results must be identical + # This test needs multiple iterations and significant batch size to simulate real + # training of a CNN of MNIST-like data. + # see https://github.com/pytorch/pytorch/issues/54771 + rref = rpc.remote(dst, MyConvNetForMNIST, args=(remote_device,)) + for _ in range(10): + rref_x = RRef(torch.randn(200, 1, 28, 28).to(local_device)) + actual = rref.remote().forward(rref_x, True).to_here() + expected = rref.rpc_sync().forward(rref_x, True) + self.assertEqual(actual, expected) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_rref_as_arg_synchronization1(self): + self._test_rref_as_arg_synchronization("cuda:0", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_as_arg_synchronization2(self): + self._test_rref_as_arg_synchronization("cuda:1", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_as_arg_synchronization3(self): + self._test_rref_as_arg_synchronization("cuda:1", "cuda:1") + + @skip_if_lt_x_gpu(2) + def test_rref_as_arg_synchronization4(self): + self._test_rref_as_arg_synchronization("cuda:0", "cuda:1") + + @skip_if_lt_x_gpu(1) + def test_rref_as_arg_synchronization5(self): + self._test_rref_as_arg_synchronization( + "cuda:0", + "cuda:0", + [["cuda:0"] for _ in range(4)], # devicesOptions + ) + + @staticmethod + def _rref_relay(rref): + return rref.to_here() + + def _test_rref_forward_synchronization(self, local_device, remote_device): + options = self.rpc_backend_options + + input_src = worker_name(0) + model_dst = worker_name(1) + out_relay = worker_name(2) + + if self.rank == 0: + # for 1) model construction 2) forward execution + options.set_device_map(model_dst, {local_device: remote_device}) + + # Forward output will be first copied to the relay node before + # returning to the worker. This is intentional, to test RRef + # forward CUDA stream synchronizations. + options.set_device_map(out_relay, {local_device: local_device}) + elif self.rank == 1: + # worker1 hosts the model and runs forward. The forward functions + # calls RRef.to_here(), hence needs to configure the device map + options.set_device_map(input_src, {remote_device: local_device}) + elif self.rank == 2: + # worker2 will get the out RRef and call to_here() and hence, needs + # to configure device map. + options.set_device_map(model_dst, {local_device: remote_device}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + if self.rank == 0: + # This test compares rref.rpc_sync().forward(x) vs rref.remote().forward(x).to_here() + # If to_here() is properly synchronized with forward(x) the results must be identical + # This test needs multiple iterations and significant batch size to simulate real + # training of a CNN of MNIST-like data. + # see https://github.com/pytorch/pytorch/issues/54771 + rref = rpc.remote(model_dst, MyConvNetForMNIST, args=(remote_device,)) + for _ in range(10): + rref_input = RRef(torch.randn(200, 1, 28, 28).to(local_device)) + rref_out = rref.remote().forward(rref_input, True) + out = rpc.remote( + out_relay, + TensorPipeAgentCudaRpcTest._rref_relay, + args=(rref_out,) + ).to_here() + expected = rref.rpc_sync().forward(rref_input, True) + self.assertEqual(out, expected) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_rref_forward_synchronization1(self): + self._test_rref_forward_synchronization("cuda:0", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_forward_synchronization2(self): + self._test_rref_forward_synchronization("cuda:0", "cuda:1") + + @skip_if_lt_x_gpu(2) + def test_rref_forward_synchronization3(self): + self._test_rref_forward_synchronization("cuda:1", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_forward_synchronization4(self): + self._test_rref_forward_synchronization("cuda:1", "cuda:1") + + def _test_owner_rref_forward_synchronization(self, local_device, remote_device): + if self.rank == 0: + options = self.rpc_backend_options + options.set_device_map("w0", {local_device: remote_device}) + rpc.init_rpc( + "w0", + rank=0, + world_size=1, + rpc_backend_options=options + ) + + model = rpc.remote( + "w0", torch.nn.Linear, (2048, 20000) + ).remote().to(remote_device) + for _ in range(30): + data = torch.rand(2048, 2048).to(local_device) + output = model.rpc_sync().forward(data) + # to_here() internally calls localValue as the caller is + # the owner of the RRef. + v0 = rpc.RRef(output).remote().sum().to_here().item() + v1 = output.sum().item() + self.assertEqual(v0, v1) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_owner_rref_forward_synchronization1(self): + self._test_owner_rref_forward_synchronization("cuda:0", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_owner_rref_forward_synchronization2(self): + self._test_owner_rref_forward_synchronization("cuda:0", "cuda:1") + + @skip_if_lt_x_gpu(2) + def test_owner_rref_forward_synchronization3(self): + self._test_owner_rref_forward_synchronization("cuda:1", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_owner_rref_forward_synchronization4(self): + self._test_owner_rref_forward_synchronization("cuda:1", "cuda:1") + + @staticmethod + def _return_tensor_view(i): + x = torch.ones(1000, 200).cuda(0) * i + torch.cuda._sleep(10 * FIFTY_MIL_CYCLES) + # serialization of the return value will create a new tensor from the + # view, which is done outside of the user function. + return x.split(100)[0] + + @skip_if_lt_x_gpu(1) + def test_tensor_view_as_return_value(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {0 : 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + futs = [ + rpc.rpc_async( + dst, + TensorPipeAgentCudaRpcTest._return_tensor_view, + args=(i,) + ) for i in range(5) + ] + + for i in range(5): + self.assertEqual(torch.ones(100, 200) * i, futs[i].wait()) + + rpc.shutdown() + + @skip_if_lt_x_gpu(2) + def test_devices_option_mismatch(self): + with self.assertRaisesRegex( + ValueError, + "Node worker0 has unexpected source devices in its device map for worker1" + ): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {0 : 0}) + options.set_devices([1]) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rpc.shutdown() + + @skip_if_lt_x_gpu(2) + def test_devices_option_mismatch_reverse(self): + with self.assertRaisesRegex( + ValueError, + "Node worker0 has unexpected target devices in its device map for worker1" + ): + dst = worker_name((self.rank + 1) % self.world_size) + + options = rpc.TensorPipeRpcBackendOptions( + init_method=self.rpc_backend_options.init_method, + num_worker_threads=self.rpc_backend_options.num_worker_threads, + device_maps={dst: {0 : 1}}, + devices=[0] + ) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_cuda_future_device_as_int(self): + Future(devices=[0]) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_device_as_str(self): + Future(devices=["cuda:0"]) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_device_as_device(self): + Future(devices=[torch.device("cuda", 0)]) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_device_not_cuda(self): + with self.assertRaisesRegex( + ValueError, "Expected devices to have indices, got cpu" + ): + Future(devices=["cpu"]) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_cuda_tensor(self): + self._test_cuda_future_extraction( + wrapper=lambda t: t, unwrapper=lambda v: v, sparse_tensor=False + ) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_list_with_cuda_tensor(self): + self._test_cuda_future_extraction( + wrapper=lambda t: [t], unwrapper=operator.itemgetter(0), sparse_tensor=False + ) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_custom_class_with_cuda_tensor(self): + self._test_cuda_future_extraction( + wrapper=TensorWrapper, unwrapper=lambda v: v.tensor, sparse_tensor=False + ) + + @skip_if_lt_x_gpu(2) + def test_cuda_future_callback_changes_devices(self): + # We check proper CUDA stream synchronization by filling the tensor with + # the expected value in one stream, and reading it from another stream. + tensor0 = torch.zeros((100,), device="cuda:0") + tensor1 = torch.zeros((100,), device="cuda:1") + parent_future = Future(devices=["cuda:0", "cuda:1"]) + + def cb(fut): + t0 = fut.value() + tensor1.copy_(t0, non_blocking=True) + return tensor1 + + child_future = parent_future.then(cb) + with torch.cuda.device("cuda:0"): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + tensor0.fill_(1) + parent_future.set_result(tensor0) + with torch.cuda.device("cuda:1"): + another_stream = torch.cuda.Stream() + with torch.cuda.stream(another_stream): + self.assertTrue(torch.eq(child_future.wait(), 1).all().item()) + + @skip_if_lt_x_gpu(2) + def test_cuda_future_value_on_bad_device(self): + tensor0 = torch.zeros((100,), device="cuda:0") + tensor1 = torch.zeros((100,), device="cuda:1") + parent_future = Future(devices=["cuda:1"]) + + # As a plus, we test that futures still invoke callbacks even in case of + # error, and that the child futures are successful if those callbacks + # don't access the parent future. + def cb(fut): + with torch.cuda.device("cuda:1"): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + tensor1.fill_(1) + return tensor1 + + child_future = parent_future.then(cb) + with torch.cuda.device("cuda:0"): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + tensor0.fill_(1) + parent_future.set_result(tensor0) + with self.assertRaisesRegex( + ValueError, + r"The result contained tensors residing on device\(s\) cuda:0 " + r"which are not among the expected device\(s\) cuda:1", + ): + parent_future.wait() + with torch.cuda.device("cuda:1"): + another_stream = torch.cuda.Stream() + with torch.cuda.stream(another_stream): + self.assertTrue(torch.eq(child_future.wait(), 1).all().item()) + + @skip_if_lt_x_gpu(1) + def test_async_execution_with_cuda_future(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {"cuda:0": "cuda:0"}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + t = torch.zeros((100,), device="cuda:0") + fut = rpc.rpc_async(dst, async_cuda_sleep_and_set_to_one, args=(t,)) + another_stream = torch.cuda.Stream("cuda:0") + with torch.cuda.stream(another_stream): + self.assertTrue(torch.eq(fut.wait(), 1).all().item()) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_async_execution_nested_with_cuda_future(self): + dst = worker_name((self.rank + 1) % self.world_size) + nested_dst = worker_name((self.rank + 2) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {"cuda:0": "cuda:0"}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + a = torch.ones((100,), device="cuda:0") + b = torch.ones((100,), device="cuda:0") + c = torch.ones((100,), device="cuda:0") + fut = rpc.rpc_async(dst, async_cuda_nested_add, args=(nested_dst, a, b, c)) + another_stream = torch.cuda.Stream("cuda:0") + with torch.cuda.stream(another_stream): + self.assertTrue(torch.eq(fut.wait(), 3).all().item()) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_cuda_future_modify_tensor_inplace(self): + tensor = torch.zeros((100,), device="cuda:0") + future = Future(devices=["cuda:0"]) + future.set_result(tensor) + # It's weird to modify the value of a future once it's complete, but + # technically possible. Currently this is considered undefined behavior + # (in practice the future will ignore the modification and still + # synchronize with the original value). We could one day add logic to + # detect and warn or throw in such cases, but for now we just check that + # this doesn't crash. + tensor.fill_(1) + future.wait() + + @skip_if_lt_x_gpu(1) + def test_cuda_future_replace_tensor(self): + tensor_list = [torch.zeros((100,), device="cuda:0")] + future = Future(devices=["cuda:0"]) + future.set_result(tensor_list) + # It's weird to modify the value of a future once it's complete, but + # technically possible. Currently this is considered undefined behavior + # (in practice the future will ignore the modification and still + # synchronize with the original value). We could one day add logic to + # detect and warn or throw in such cases, but for now we just check that + # this doesn't crash. + # We set things up so that the original tensor contained in the list + # gets deleted once we replace it with the other one. This will + # invalidate any cached information held by the future. + tensor_list[0] = torch.ones((100,), device="cuda:0") + future.wait() + + @skip_if_lt_x_gpu(1) + def test_rref_with_unpickleable_attributes(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {"cuda:0": "cuda:0"}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rref = rpc.remote(dst, TensorWrapper, args=(torch.zeros(42, device="cuda:0"),)) + rref.rpc_sync().increase(1) + ret = rref.rpc_sync().sum() + self.assertEqual(ret, 42) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_cuda_sparse_tensor(self): + self._test_cuda_future_extraction( + wrapper=lambda t: t, unwrapper=lambda v: v, sparse_tensor=True + ) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_list_with_cuda_sparse_tensor(self): + self._test_cuda_future_extraction( + wrapper=lambda t: [t], unwrapper=operator.itemgetter(0), sparse_tensor=True + ) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_custom_class_with_cuda_sparse_tensor(self): + self._test_cuda_future_extraction( + wrapper=TensorWrapper, unwrapper=lambda v: v.tensor, sparse_tensor=True + ) diff --git a/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/tensorpipe_rpc_agent_test_fixture.py b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/tensorpipe_rpc_agent_test_fixture.py new file mode 100644 index 0000000000000000000000000000000000000000..629b0328f1aed6d14d71589faee40a9b702a1b22 --- /dev/null +++ b/lib/python3.12/site-packages/torch/testing/_internal/distributed/rpc/tensorpipe_rpc_agent_test_fixture.py @@ -0,0 +1,34 @@ +# mypy: allow-untyped-defs + +import torch.distributed.rpc as rpc +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) +from torch.testing._internal.common_distributed import ( + tp_transports, +) + + +class TensorPipeRpcAgentTestFixture(RpcAgentTestFixture): + @property + def rpc_backend(self): + return rpc.backend_registry.BackendType[ + "TENSORPIPE" + ] + + @property + def rpc_backend_options(self): + return rpc.backend_registry.construct_rpc_backend_options( + self.rpc_backend, + init_method=self.init_method, + _transports=tp_transports() + ) + + def get_shutdown_error_regex(self): + # FIXME Once we consolidate the error messages returned by the + # TensorPipe agent put some more specific regex here. + error_regexes = [".*"] + return "|".join([f"({error_str})" for error_str in error_regexes]) + + def get_timeout_error_regex(self): + return "RPC ran for more than"