Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,801 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import _operator
|
||||
import itertools
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
|
||||
from torch.fx import Node
|
||||
from torch.fx._compatibility import compatibility
|
||||
from torch.multiprocessing.reductions import StorageWeakRef
|
||||
from torch.utils import _pytree as pytree
|
||||
from torch.utils._pytree import tree_map_only
|
||||
|
||||
|
||||
__all__ = ["reinplace"]
|
||||
|
||||
|
||||
class _ViewType(Enum):
|
||||
NonView = 0
|
||||
SingleOutputView = 1
|
||||
MultiOutputView = 2
|
||||
|
||||
|
||||
def _is_view_op(tgt):
|
||||
if tgt is not None and isinstance(tgt, torch._ops.OpOverload):
|
||||
schema = tgt._schema
|
||||
if len(schema.arguments) > 0:
|
||||
first_arg = schema.arguments[0]
|
||||
# check if op is a view
|
||||
return (
|
||||
first_arg.alias_info is not None and not first_arg.alias_info.is_write
|
||||
)
|
||||
|
||||
|
||||
def _get_view_type(tgt) -> _ViewType:
|
||||
if tgt is not None and isinstance(tgt, torch._ops.OpOverload):
|
||||
schema = tgt._schema
|
||||
if len(schema.arguments) > 0:
|
||||
first_arg = schema.arguments[0]
|
||||
# check if op is a view
|
||||
if first_arg.alias_info is not None and not first_arg.alias_info.is_write:
|
||||
# check if op is a multi-output view
|
||||
if "*" in first_arg.alias_info.after_set:
|
||||
return _ViewType.MultiOutputView
|
||||
else:
|
||||
return _ViewType.SingleOutputView
|
||||
return _ViewType.NonView
|
||||
|
||||
|
||||
# Stores a bunch of metadata related to functionalization each node.
|
||||
# Relevant metadata:
|
||||
# n.meta['fake_result']: FakeTensor (same type as the output of the node, but with FakeTenors instead of Tensors)
|
||||
# The fake tensor output from running the current node
|
||||
# n.meta['view_of']: Node
|
||||
# If the current node n is a view of some base tensor, the 'view_of' field tells us which
|
||||
# view node was used to generate the current node (a view tensor).
|
||||
# This information actually makes `fake_result` redundant, but we can use `fake_result`
|
||||
# to sanity check that our aliasing information is correct.
|
||||
@compatibility(is_backward_compatible=False)
|
||||
class _FunctionalizationMetadataProp(torch.fx.Interpreter):
|
||||
def run_node(self, node: Node):
|
||||
self.node_counter += 1
|
||||
result = super().run_node(node)
|
||||
node.meta["fake_result"] = result
|
||||
node.meta["node_idx"] = self.node_counter
|
||||
|
||||
# (1) Update metadata with the list of nodes that are used by this node
|
||||
# copy_() doesn't read from its first argument; it writes to it, overwriting previous data.
|
||||
# We don't want to treat it as "being used as an input".
|
||||
node_args = node.args
|
||||
if node.target is torch.ops.aten.copy_.default:
|
||||
node_args = node_args[1:]
|
||||
|
||||
# (2) Update metadata to track aliasing information about view tensor nodes.
|
||||
if node.op == "call_function":
|
||||
view_type = _get_view_type(node.target)
|
||||
if view_type == _ViewType.SingleOutputView:
|
||||
if not isinstance(node.args[0], Node):
|
||||
raise AssertionError(f"Expected Node, got {type(node.args[0])}")
|
||||
node.meta["view_of"] = node.args[0]
|
||||
elif view_type == _ViewType.MultiOutputView:
|
||||
self.multi_output_view_nodes[node] = node.args[0]
|
||||
|
||||
# Check if we returned a multi-output view,
|
||||
# and we're now grabbing the individual views from the output.
|
||||
#
|
||||
# For multi-output views, we want to map each output view to the base,
|
||||
# but this mapping involves two separate nodes in FX IR.
|
||||
# e.g. "a, b = x_1.split(...)" becomes:
|
||||
# %split_tensor : [num_users=2] = call_function[target=torch.ops.aten.split.Tensor](args = (%x_1, 2), kwargs = {})
|
||||
# %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%split_tensor, 0), kwargs = {})
|
||||
# %getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%split_tensor, 1), kwargs = {})
|
||||
# And we'd like to set:
|
||||
# getitem1.meta['view_of'] = x_1
|
||||
elif node.target is _operator.getitem:
|
||||
list_arg = node.args[0]
|
||||
maybe_base_of_view = self.multi_output_view_nodes.get(list_arg, None)
|
||||
if maybe_base_of_view is not None:
|
||||
# Note: we could also track indexing info here for multi-output views.
|
||||
# I don't think this metadata is strictly needed for de-functionalization.
|
||||
if not isinstance(maybe_base_of_view, Node):
|
||||
raise AssertionError(
|
||||
f"Expected Node, got {type(maybe_base_of_view)}"
|
||||
)
|
||||
node.meta["view_of"] = maybe_base_of_view
|
||||
|
||||
if "view_of" in node.meta:
|
||||
# We're linking the current node with its first argument as views.
|
||||
# Assert here that this is actually the case, and their storages are the same.
|
||||
if not isinstance(node.meta["fake_result"], FakeTensor):
|
||||
raise AssertionError("Expected FakeTensor in fake_result")
|
||||
if not isinstance(node.meta["view_of"].meta["fake_result"], FakeTensor):
|
||||
raise AssertionError("Expected FakeTensor in view_of fake_result")
|
||||
view_storage = StorageWeakRef(node.meta["fake_result"]._typed_storage())
|
||||
base_storage = StorageWeakRef(
|
||||
node.meta["view_of"].meta["fake_result"]._typed_storage()
|
||||
)
|
||||
if view_storage != base_storage:
|
||||
raise AssertionError("view_storage != base_storage")
|
||||
return result
|
||||
|
||||
def propagate(self, *args):
|
||||
self.multi_output_view_nodes = {}
|
||||
self.node_counter = -1
|
||||
|
||||
with FakeTensorMode() as mode:
|
||||
fake_args = [
|
||||
mode.from_tensor(a) if isinstance(a, torch.Tensor) else a for a in args
|
||||
]
|
||||
return super().run(*fake_args)
|
||||
|
||||
|
||||
def _schemas_match(functional_schema, inplace_schema):
|
||||
names_match = (
|
||||
inplace_schema.name.endswith("_")
|
||||
and inplace_schema.name[:-1] == functional_schema.name
|
||||
)
|
||||
arg_types_match = len(functional_schema.arguments) == len(
|
||||
inplace_schema.arguments
|
||||
) and all(
|
||||
a1.type == a2.type
|
||||
for a1, a2 in zip(functional_schema.arguments, inplace_schema.arguments)
|
||||
)
|
||||
# for the inplace op, its first argument should be mutable
|
||||
if not (
|
||||
inplace_schema.arguments[0].alias_info is not None
|
||||
and inplace_schema.arguments[0].alias_info.is_write
|
||||
):
|
||||
raise AssertionError("First argument of inplace op must be mutable")
|
||||
# and its remaining arguments shouldn't be.
|
||||
if not all(a.alias_info is None for a in inplace_schema.arguments[1:]):
|
||||
raise AssertionError("Remaining arguments of inplace op must not be mutable")
|
||||
return names_match and arg_types_match
|
||||
|
||||
|
||||
# TODO: this should be beefed up to be able to properly re-inplace with:
|
||||
# - mutating ops (e.g. _fused_moving_avg_obs_fq_helper)
|
||||
# - out= ops (e.g. angle -> angle.out)
|
||||
# TODO: we should also figure this info out using torchgen.
|
||||
def _maybe_get_inplace_op(op):
|
||||
# __module__ seems broken; it returns torch._ops.aten which doesn't exist
|
||||
if not isinstance(op, torch._ops.OpOverload):
|
||||
return None
|
||||
# Some view ops have inplace variants (as_strided_, etc),
|
||||
# but we do NOT want the reinplacing pass to directly add these into the program.
|
||||
# (they'll require extra special handling, aren't aren't really useful for perf anyway)
|
||||
if _is_view_op(op):
|
||||
return None
|
||||
op_namespace = op.__module__.split(".")[-1]
|
||||
op_base_name = op.overloadpacket.__name__
|
||||
maybe_namespace_module = getattr(torch.ops, op_namespace)
|
||||
maybe_inplace_op = (
|
||||
None
|
||||
if maybe_namespace_module is None
|
||||
else getattr(maybe_namespace_module, f"{op_base_name}_", None)
|
||||
)
|
||||
if maybe_inplace_op is None:
|
||||
return None
|
||||
|
||||
inplace_overloads = [
|
||||
getattr(maybe_inplace_op, overload_name)
|
||||
for overload_name in maybe_inplace_op.overloads()
|
||||
]
|
||||
inplace_overloads_with_matching_schemas = [
|
||||
f for f in inplace_overloads if _schemas_match(op._schema, f._schema)
|
||||
]
|
||||
# Just because foo() and foo_() are both existing operators,
|
||||
# They aren't guaranteed to have compatible schemas.
|
||||
# For example, pow.Scalar(Scalar self, Tensor exponent) has no valid inplace variant,
|
||||
# Even though several overloads of pow_ exist.
|
||||
if len(inplace_overloads_with_matching_schemas) == 0:
|
||||
return None
|
||||
if len(inplace_overloads_with_matching_schemas) != 1:
|
||||
raise AssertionError(
|
||||
f"Expected exactly 1 matching inplace overload, got "
|
||||
f"{len(inplace_overloads_with_matching_schemas)}"
|
||||
)
|
||||
inplace_op = inplace_overloads_with_matching_schemas[0]
|
||||
return inplace_op
|
||||
|
||||
|
||||
_VIEW_INVERSE_MAP: dict[Callable[..., Any], Callable[..., Any]] = {
|
||||
torch.ops.aten.diagonal_scatter.default: torch.ops.aten.diagonal.default,
|
||||
torch.ops.aten.select_scatter.default: torch.ops.aten.select.int,
|
||||
torch.ops.aten.slice_scatter.default: torch.ops.aten.slice.Tensor,
|
||||
torch.ops.aten.as_strided_scatter.default: torch.ops.aten.as_strided.default,
|
||||
}
|
||||
|
||||
|
||||
# This function, given a set of set of (aliased) tensor nodes,
|
||||
# Returns any nodes in the graph that *use* any of the aliases, that occur *after* op_index
|
||||
# in the node ordering.
|
||||
def _get_all_later_node_usages(tensor_aliases: set[Node], op_index: int):
|
||||
def _add_if_tensor(x, set_):
|
||||
if isinstance(x, FakeTensor):
|
||||
set_.add(StorageWeakRef(x._typed_storage()))
|
||||
|
||||
nodes_used_after = set()
|
||||
for t in tensor_aliases:
|
||||
# get all nodes that use the current alias
|
||||
usage_nodes = t.users
|
||||
for n in usage_nodes:
|
||||
# We only care about usages after the current node
|
||||
if "node_idx" not in n.meta or n.meta["node_idx"] <= op_index:
|
||||
continue
|
||||
# We also don't care about intermediate view ops.
|
||||
# They only matter if their output is then used elsewhere
|
||||
# (either in an out-of-place op, or as an output to the function).
|
||||
if n in tensor_aliases:
|
||||
if (
|
||||
isinstance(n.target, torch._ops.OpOverload)
|
||||
or n.target is _operator.getitem
|
||||
):
|
||||
continue
|
||||
nodes_used_after.add(n)
|
||||
return nodes_used_after
|
||||
|
||||
|
||||
# Given an op that we're trying to re-inplace, "b = foo(a)",
|
||||
# And given a {view}_scatter op that shows up later in the graph, "y = {view}_scatter(base, x, args...)"
|
||||
# Then re-inplacing `foo()` would allow us to remove the `{view}_scatter` op entirely, IF:
|
||||
# If there are any aliases in the alias_set(a) that satisfy:
|
||||
# (1) The base of "alias", "alias_base", has the same size/stride/offset metadata as "base"
|
||||
# (2) The output of running {view}(alias, args...) gives you the same size/stride/offset metadata
|
||||
# as "alias"
|
||||
def _get_view_inverse_node_usages(
|
||||
later_node_usages: set[Node], self_aliases: set[Node]
|
||||
) -> set[Node]:
|
||||
def matching_view_metadata(a, b):
|
||||
return (
|
||||
a.size() == b.size()
|
||||
and a.stride() == b.stride()
|
||||
and a.storage_offset() == b.storage_offset()
|
||||
)
|
||||
|
||||
view_inverse_nodes = set()
|
||||
# Go through them in node order, so we can see chains of view_scatter ops.
|
||||
for n in sorted(later_node_usages, key=lambda x: x.meta["node_idx"]):
|
||||
if n.target not in _VIEW_INVERSE_MAP:
|
||||
continue
|
||||
base = n.args[0]
|
||||
mutated_view = n.args[1]
|
||||
if not isinstance(base, Node):
|
||||
raise AssertionError(f"Expected Node for base, got {type(base)}")
|
||||
if not isinstance(base.meta["fake_result"], FakeTensor):
|
||||
raise AssertionError("Expected FakeTensor in base.meta['fake_result']")
|
||||
if not isinstance(mutated_view, Node):
|
||||
raise AssertionError(
|
||||
f"Expected Node for mutated_view, got {type(mutated_view)}"
|
||||
)
|
||||
if not isinstance(mutated_view.meta["fake_result"], FakeTensor):
|
||||
raise AssertionError(
|
||||
"Expected FakeTensor in mutated_view.meta['fake_result']"
|
||||
)
|
||||
if isinstance(n.target, str):
|
||||
raise AssertionError("n.target should not be a string")
|
||||
# Check that this view_inverse op actually corresponds to taking doing the inverse
|
||||
# of one of our existing self_alias nodes.
|
||||
original_view = _VIEW_INVERSE_MAP[n.target]
|
||||
for self_alias in self_aliases:
|
||||
# We're looking for some alias of the self arg, "alias",
|
||||
# that was created from some op `alias = foo(base, args...)`
|
||||
# such that the current _scatter op "inverts" that foo call.
|
||||
# We can check that by running the original op again, and checking that the strides match.
|
||||
if "view_of" not in self_alias.meta:
|
||||
continue
|
||||
self_alias_base = self_alias.meta["view_of"]
|
||||
try:
|
||||
# The we're trying to reuse the args from the view_scatter call inside of the corresponding
|
||||
# view op, which might throw. This just indicates that view_scatter op isn't a valid inverse
|
||||
# of the current alias we're looking at.
|
||||
view_replay_metadata = original_view(
|
||||
self_alias_base.meta["fake_result"], *n.args[2:], **n.kwargs
|
||||
)
|
||||
expected_metadata = self_alias.meta["fake_result"]
|
||||
# If the alias and its base both have matching metadata, then this view_scatter op is valid to re-inplace.
|
||||
if matching_view_metadata(
|
||||
self_alias_base.meta["fake_result"], base.meta["fake_result"]
|
||||
) and matching_view_metadata(view_replay_metadata, expected_metadata):
|
||||
view_inverse_nodes.add(n)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return view_inverse_nodes
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=True)
|
||||
def reinplace(gm, *sample_args):
|
||||
r"""
|
||||
Given an fx.GraphModule, modifies it to perform "reinplacing",
|
||||
mutating the nodes of the graph.
|
||||
We look for out-of-place op call sites like ``b = a.add(...)``,
|
||||
and convert them to be inplace (``b = a.add_(...)``),
|
||||
as long as the input to the current operator ("a") isn't reused
|
||||
anywhere later in the graph.
|
||||
|
||||
This pass currently expects to operate on a **functional, ATen** graph.
|
||||
This can be obtained by running ``make_fx(functionalize(f))``.
|
||||
|
||||
Sample inputs are needed to determine aliasing relationships of the inputs.
|
||||
In general, we can't reinplace node ``b = a.add(...)`` if "a" aliases any of the
|
||||
inputs to the program.
|
||||
|
||||
Given a node ``b = foo(a, args...)`` the algorithm for re-inplacing is as follows:
|
||||
|
||||
**(1)** Perform some initial checks on the metadata of "a" and "args..."
|
||||
that can disqualify them from being reinplaced.
|
||||
|
||||
- **(1a)** Check that the self argument we're attempting to reinplace
|
||||
has acceptable dtype/size metadata to reinplace with.
|
||||
|
||||
For example, if we have::
|
||||
|
||||
a = torch.ones(1)
|
||||
b = torch.ones(10)
|
||||
out = torch.add(a, b)
|
||||
|
||||
We can't turn that into ``a.add_(b)`` because that would require resizing "a".
|
||||
|
||||
Similarly, we can't convert ``torch.ge(a, b)`` into ``a.ge_(b)``,
|
||||
because that would require changing a's dtype (from e.g. float32 to bool).
|
||||
Note that in this specific example, we could technically do better..
|
||||
|
||||
If we see the pattern::
|
||||
|
||||
a_1 = a.ge(b)
|
||||
a_2 = aten._to_copy(a_1, a.dtype)
|
||||
|
||||
Then this should be valid to completely re-inplace
|
||||
(this is exactly what functionalization will emit when it sees ``a.ge_(b)``).
|
||||
|
||||
This optimization is only really important for user programs
|
||||
that directly use inplace comparison ops though.
|
||||
|
||||
We also cannot re-inplace on tensors that have overlapping memory,
|
||||
e.g. ``torch.ones(1).expand(4, 4).add_(1)``.
|
||||
|
||||
- **(1b)** Check if "a" is an alias of any of the program inputs.
|
||||
|
||||
If it is, skip and move to the next node.
|
||||
Inplace'ing an op that would cause it to mutate a program is not sound,
|
||||
because that would be a side effect visible to the user.
|
||||
|
||||
NOTE: there's a future optimization that we should make:
|
||||
if "a" is a (alias of a) program input, but later in the program
|
||||
there is a node that looks like ``a.copy_(...)``,
|
||||
then re-inplacing is ok to do - we are temporarily reusing a's buffer,
|
||||
which will later be overwritten by the ``copy_()`` call.
|
||||
|
||||
This will be an important optimization to have for programs that mutate
|
||||
their inputs. It currently isn't implemented though.
|
||||
|
||||
- **(1c)** Check if "a" and "args..." alias.
|
||||
|
||||
For example, re-inplacing to create code like the below
|
||||
isn't guaranteed to be sound::
|
||||
|
||||
aten.mul_(a, a)
|
||||
|
||||
**(2)** Check that "a" and all of its outstanding aliases are not used anywhere
|
||||
later in the graph. If this is the case, then it's safe to re-inplace
|
||||
to ``b = foo_(a)``.
|
||||
|
||||
There are a few caveats to this, explained in more detail below:
|
||||
|
||||
- (a) If "a" is used later as an argument to a view op, that is okay.
|
||||
It's only a problem if "a" (or that view) is later passed
|
||||
into a normal operator, or if it is returned as the program output.
|
||||
- (b) If "a" is a repeat argument in ``foo()``, then don't reinplace.
|
||||
Most ATen kernels don't make any guarantees that this is sound,
|
||||
e.g. if you do ``aten.mul_(a, a)``.
|
||||
So we'll just ban re-inplacing in this case.
|
||||
- (c) If "a" is used as an input into a view "inverse" / "scatter"
|
||||
operator, it is potentially fine to re-inplace
|
||||
(and remove that scatter operator from the graph).
|
||||
See below for a more detailed example.
|
||||
|
||||
NOTE: there is an optimization in this step that is crucial
|
||||
to fully recovering performance from functionalization.
|
||||
|
||||
Given this program::
|
||||
|
||||
def f(x):
|
||||
a = torch.ops.aten.add(x, x)
|
||||
b = torch.ops.aten.diagonal(a)
|
||||
torch.ops.aten.fill_(b, 0)
|
||||
return d
|
||||
|
||||
Functionalization will emit the following::
|
||||
|
||||
def f(x):
|
||||
a = torch.ops.aten.add(x, x)
|
||||
b = torch.ops.aten.diagonal(a, 0, 1)
|
||||
b_updated = torch.ops.aten.fill(b, 0)
|
||||
a_updated = torch.ops.aten.diagonal_scatter(a, b_updated, 0, 1)
|
||||
return a_updated
|
||||
|
||||
Ordinarily, we would not be able to reinplace the fill,
|
||||
because "b" aliases with "a" which is used by the diagonal_scatter call.
|
||||
|
||||
"re-inplacing" is on the hook for figuring out that it is ok to
|
||||
completely remove the expensive diagonal_scatter call, if we re-inplace
|
||||
the add().
|
||||
|
||||
So, for every ``alias in alias_set(a)``, instead of checking
|
||||
that "alias" is not used anywhere later in the graph,
|
||||
we check that EITHER:
|
||||
|
||||
- (a) alias is not used anywhere later in the graph, OR
|
||||
- (b) alias is used exactly once later on in the graph,
|
||||
in the following op::
|
||||
|
||||
out = foo_scatter(alias, x, args...)
|
||||
|
||||
where the following must hold:
|
||||
|
||||
- (i) ``foo_scatter`` is the "inverse" operator for foo.
|
||||
This only applies to "foo" ops that are view operators,
|
||||
which view into a subset of the original tensor's memory.
|
||||
In practice, there are ~4 operators where this applies::
|
||||
|
||||
diagonal -> diagonal_scatter
|
||||
slice -> slice_scatter
|
||||
select -> select_scatter
|
||||
as_strided -> as_strided_scatter
|
||||
|
||||
- (ii) "args..." are the same between the ``foo()`` and
|
||||
``foo_scatter()`` calls.
|
||||
|
||||
**(3)** Perform the actual re-inplacing on foo!
|
||||
|
||||
(3b) is the common case, but special care is needed for
|
||||
``{view}_scatter`` (3a).
|
||||
|
||||
- **(3a)** ``{view}_scatter`` ops.
|
||||
|
||||
Consider this program::
|
||||
|
||||
a = torch.zeros(2, 2)
|
||||
b = torch.ones(2)
|
||||
a[0] = b
|
||||
|
||||
Post functionalization, that will look like::
|
||||
|
||||
a = torch.zeros(2)
|
||||
b = torch.ones(1)
|
||||
a_updated = torch.select_scatter(a, b, 0, 0)
|
||||
|
||||
In this case though, there is no "functional" op to re-inplace!
|
||||
Instead, we'd like to directly remove the select_scatter call.
|
||||
We already know from (3) that this is valid,
|
||||
because "a" has no later usages in the graph.
|
||||
|
||||
We perform the re-inplacing on the ``{view}_scatter`` op like so.
|
||||
|
||||
Before::
|
||||
|
||||
a_updated = torch.select_scatter(a, b, args...)
|
||||
|
||||
After::
|
||||
|
||||
a_slice = a.select(a, args...)
|
||||
a_slice.copy_(b)
|
||||
|
||||
- **(3b)** Otherwise, replace the functional op with its inplace variant.
|
||||
|
||||
Before::
|
||||
|
||||
b = foo(a, args...)
|
||||
|
||||
After::
|
||||
|
||||
a.foo_(args...)
|
||||
|
||||
**(4)** Finally, after converting either::
|
||||
|
||||
# Before: # After:
|
||||
b = foo(a) foo_(a)
|
||||
|
||||
or::
|
||||
|
||||
# Before:
|
||||
b = {slice}_scatter(a, mutated_slice, args...)
|
||||
# After:
|
||||
slice = {slice}(a, args...)
|
||||
slice.copy_(mutated_slice)
|
||||
|
||||
We now need to find all later nodes that use "b" as an argument
|
||||
and update them to take in "a" instead.
|
||||
|
||||
Note that for the majority of inplace ops, this isn't actually necessary
|
||||
(because most inplace ops return "self" as their output).
|
||||
This isn't generally true for all mutable ops though, which is why
|
||||
we need to actually replace all of the arguments.
|
||||
|
||||
We also need to update our metadata of ``Dict[StorageWeakRef, Set[Node]]``,
|
||||
that maps a given tensor storage to the set of all nodes that take in that
|
||||
storage as an input.
|
||||
Specifically, re-inplacing ``b = foo(a)`` causes "a" and "b"'s sets to get
|
||||
fused together.
|
||||
|
||||
**(5)** Any ``view_inverse/scatter`` nodes that were identified as
|
||||
"it's ok to ignore them" during step (3) get manually deleted from the graph.
|
||||
Their outputs are no longer used, so technically standard DCE would be able
|
||||
to do this, but we can no longer run FX's DCE pass now that we have mutable
|
||||
ops in the graph.
|
||||
"""
|
||||
_FunctionalizationMetadataProp(gm).propagate(*sample_args)
|
||||
|
||||
# Useful debug printing
|
||||
# def _print(x):
|
||||
# if isinstance(x, FakeTensor):
|
||||
# print(f'fake_result: {StorageWeakRef(x._typed_storage()).cdata}')
|
||||
|
||||
# for n in gm.graph.nodes:
|
||||
# print(n.format_node())
|
||||
# if hasattr(n, 'meta'):
|
||||
# print(f'node_idx: {n.meta["node_idx"]}')
|
||||
# if 'fake_result' in n.meta:
|
||||
# tree_map(_print, n.meta['fake_result'])
|
||||
# if 'view_of' in n.meta:
|
||||
# print(f'view_of: {str(n.meta["view_of"])}')
|
||||
# print()
|
||||
|
||||
# We need to know which nodes correspond to inputs (or their aliases)
|
||||
# so we know not to re-inplace them.
|
||||
# NOTE: later, we'll need to add an optimization for fully recovering performance
|
||||
# on programs that mutate inputs.
|
||||
input_storages = {
|
||||
StorageWeakRef(node.meta["fake_result"]._typed_storage())
|
||||
for node in gm.graph.nodes
|
||||
if (
|
||||
node.op == "placeholder"
|
||||
and isinstance(node.meta["fake_result"], torch.Tensor)
|
||||
)
|
||||
}
|
||||
|
||||
# We also need to know for a given node, what are all of its aliasing nodes.
|
||||
storage_to_nodes: dict[StorageWeakRef, set[Node]] = defaultdict(set)
|
||||
for n in gm.graph.nodes:
|
||||
if "fake_result" in n.meta:
|
||||
# Tree-mapping because some ops can return lists of tensors.
|
||||
def _add_to_map(x):
|
||||
if isinstance(x, FakeTensor):
|
||||
storage_to_nodes[StorageWeakRef(x._typed_storage())].add(n)
|
||||
|
||||
pytree.tree_map_(_add_to_map, n.meta["fake_result"])
|
||||
|
||||
# inplace-ify functional ops, subject to the constraints written below.
|
||||
all_later_view_inverse_nodes_to_delete = set()
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "call_function":
|
||||
# Today, the re-inplace pass on directly acts on:
|
||||
# - functional ops with an inplace variant
|
||||
# - {view}_scatter ops that can be potentially removed from the graph.
|
||||
# Both of these ops take in tensor first args, so filtering on this condition
|
||||
# makes the later code simpler.
|
||||
# We should revisit this at some point though, particularly when we also want
|
||||
# the reinplacer to be able to handle out= and mutable operators
|
||||
# and tensorlist first args (like `_foreach_` ops).
|
||||
if not isinstance(node.target, torch._ops.OpOverload):
|
||||
continue
|
||||
if len(node.target._schema.arguments) < 1:
|
||||
continue
|
||||
if type(node.target._schema.arguments[0].type) is not torch.TensorType:
|
||||
continue
|
||||
|
||||
# Step 1a: Check that the self argument we're attempting to reinplace
|
||||
# has the same size/stride as the output.
|
||||
# For example, we shouldn't try to reinplace torch.add(scalar_tensor, larger_tensor)
|
||||
# As it would require resizing scalar_tensor.
|
||||
# (We could potentially swizzle this into larger_tensor.add_(scalar_tensor),
|
||||
# this is probably an optimization to revisit later).
|
||||
self_arg = node.args[0]
|
||||
self_flattened = pytree.tree_leaves(self_arg.meta["fake_result"])
|
||||
node_flattened = pytree.tree_leaves(node.meta["fake_result"])
|
||||
self_has_wrong_metadata = False
|
||||
if len(self_flattened) == len(node_flattened):
|
||||
for self_meta, node_meta in zip(self_flattened, node_flattened):
|
||||
if self_meta.numel() != node_meta.numel():
|
||||
self_has_wrong_metadata = True
|
||||
if self_meta.dtype != node_meta.dtype:
|
||||
self_has_wrong_metadata = True
|
||||
# We also cannot re-inplace on tensors that have internal memory overlap.
|
||||
# e.g. torch.ones(1).expand(4, 4).add_(1)
|
||||
if torch._debug_has_internal_overlap(self_meta) == 1:
|
||||
self_has_wrong_metadata = True
|
||||
# Here, we (optimistically) assume that a.resize(b) is valid to re-inplace,
|
||||
# Since users should never really be calling the functional "torch.ops.aten.resize"
|
||||
# op directly in their programs.
|
||||
if self_has_wrong_metadata and node.target != torch.ops.aten.resize.default:
|
||||
continue
|
||||
|
||||
# Step 1b: ensure that the op we're trying to re-inplace isn't a program input
|
||||
self_arg_storage = StorageWeakRef(
|
||||
self_arg.meta["fake_result"]._typed_storage()
|
||||
)
|
||||
if self_arg_storage in input_storages:
|
||||
# TODO: later, add the optimization for handling `copy_()` calls in the graph.
|
||||
continue
|
||||
if len([x for x in node.args if x is self_arg]) > 1:
|
||||
# Step 1c:
|
||||
# Calling stuff like aten.mul_(a, a) isn't guaranteed to be sound,
|
||||
# so we prevent re-inplacing in this case.
|
||||
continue
|
||||
|
||||
self_arg_storage = StorageWeakRef(
|
||||
self_arg.meta["fake_result"]._typed_storage()
|
||||
)
|
||||
self_aliases = storage_to_nodes[self_arg_storage]
|
||||
|
||||
# First, we find all later usages of any of the aliases of self_arg.
|
||||
later_node_usages = _get_all_later_node_usages(
|
||||
self_aliases, node.meta["node_idx"]
|
||||
)
|
||||
# Then, we check if any of those later usages are actually view_scatter ops
|
||||
# that are safe to fully remove.
|
||||
later_view_inverse_node_usages = _get_view_inverse_node_usages(
|
||||
later_node_usages, self_aliases
|
||||
)
|
||||
|
||||
# Step 2: Check to see if the input to the op is reused later in the graph.
|
||||
# If not (same goes for its aliases), then this op is safe to re-in place.
|
||||
# This is a slightly roundabout way to check that there are no later usages of the current self argument.
|
||||
# (later_view_inverse_node_usages corresponds to "view_scatter" nodes that we are allowed to delete)
|
||||
can_reinplace = len(later_node_usages - later_view_inverse_node_usages) == 0
|
||||
if not can_reinplace:
|
||||
continue
|
||||
|
||||
# Step 3a: Special handling for when we see *_scatter operators.
|
||||
# When we see an operator like `b = torch.slice_scatter(a, ...)`,
|
||||
# instead of trying to "inplace" it into a.slice_scatter_(..._),
|
||||
# we would prefer to remove it from the graph entirely,
|
||||
# and instead copy_() the slice directly into the larger tensor.
|
||||
# See the description of the algorithm for a full example.
|
||||
if (
|
||||
node.target in _VIEW_INVERSE_MAP
|
||||
and node not in all_later_view_inverse_nodes_to_delete
|
||||
):
|
||||
view_op = _VIEW_INVERSE_MAP[node.target]
|
||||
# Before:
|
||||
# base_updated = torch.ops.aten.slice_scatter.default(base, mutated_slice, args...)
|
||||
# After:
|
||||
# slice = torch.ops.aten.slice.default(base, args...)
|
||||
# slice.copy_(mutated_slice)
|
||||
with gm.graph.inserting_before(node):
|
||||
mutated_slice_node = node.args[1]
|
||||
remaining_slice_args = node.args[2:]
|
||||
slice_node = gm.graph.create_node(
|
||||
"call_function",
|
||||
view_op,
|
||||
(self_arg,) + tuple(remaining_slice_args),
|
||||
node.kwargs,
|
||||
)
|
||||
gm.graph.create_node(
|
||||
"call_function",
|
||||
torch.ops.aten.copy_.default,
|
||||
(
|
||||
slice_node,
|
||||
mutated_slice_node,
|
||||
),
|
||||
{},
|
||||
)
|
||||
# Add the slice_scatter node to our "nodes to delete" list.
|
||||
all_later_view_inverse_nodes_to_delete.add(node)
|
||||
|
||||
else:
|
||||
# Step 3b: Check to see if this operator has an inplace variant.
|
||||
maybe_inplace_op = _maybe_get_inplace_op(node.target)
|
||||
if maybe_inplace_op is None:
|
||||
continue
|
||||
# And if so, replace it with its inplace variant.
|
||||
node.target = maybe_inplace_op
|
||||
|
||||
# At this point, 'storage_to_nodes' will be stale.
|
||||
# Now that we're inplacing `b = foo(a)`, we need to effectively
|
||||
# union together the dict values for b and a's storage.
|
||||
# Hmm... morally I think we also want to keep the `fake_result` metadata
|
||||
# up to date here, but I'm not sure how easy it is to do.
|
||||
# Maybe it's fine to wait until the end of the pass to update it.
|
||||
curr_node_storage = StorageWeakRef(
|
||||
node.meta["fake_result"]._typed_storage()
|
||||
)
|
||||
storage_to_nodes[self_arg_storage].update(
|
||||
storage_to_nodes[curr_node_storage]
|
||||
)
|
||||
storage_to_nodes[curr_node_storage].update(
|
||||
storage_to_nodes[self_arg_storage]
|
||||
)
|
||||
|
||||
# Need to remember the view_scatter view nodes we found so we can remove them alter.
|
||||
all_later_view_inverse_nodes_to_delete.update(
|
||||
later_view_inverse_node_usages
|
||||
)
|
||||
|
||||
# Step 4:
|
||||
# Now that we've replaced b = a.foo() with a.foo_(),
|
||||
# We need to replace any later usages of "b" with "a"
|
||||
for old in itertools.chain([node], later_view_inverse_node_usages):
|
||||
new = old.args[0]
|
||||
nodes_to_update = [
|
||||
n for n in old.users if n.meta["node_idx"] > node.meta["node_idx"]
|
||||
]
|
||||
for node_to_update in nodes_to_update:
|
||||
|
||||
def replace_arg(a):
|
||||
if a == old:
|
||||
return new
|
||||
return a
|
||||
|
||||
# First, replace usages of "b" with "a"
|
||||
node_to_update.args = tree_map_only(
|
||||
Node, replace_arg, node_to_update.args
|
||||
)
|
||||
node_to_update.kwargs = tree_map_only(
|
||||
Node, replace_arg, node_to_update.kwargs
|
||||
)
|
||||
|
||||
# Second, update our storage_to_nodes data structure.
|
||||
old_flattened_res = pytree.tree_leaves(old.meta["fake_result"])
|
||||
node_flattened_res = pytree.tree_leaves(
|
||||
node_to_update.meta["fake_result"]
|
||||
)
|
||||
|
||||
old_res_storage = {
|
||||
StorageWeakRef(x._typed_storage())
|
||||
for x in old_flattened_res
|
||||
if isinstance(x, FakeTensor)
|
||||
}
|
||||
node_res_storage = {
|
||||
StorageWeakRef(x._typed_storage())
|
||||
for x in node_flattened_res
|
||||
if isinstance(x, FakeTensor)
|
||||
}
|
||||
|
||||
# This will happen if we're updating a view op, e.g.
|
||||
# e.g. replacing
|
||||
# x = view(old)
|
||||
# x = view(new)
|
||||
# When that happens, we need to make sure to keep our
|
||||
# storage mapping up to date.
|
||||
#
|
||||
# We're checking for len(...) == 1 here because all view ops are guaranteed to return either a single tensor,
|
||||
# or multiple tensors that all share the same storage.
|
||||
# We can't just check equality because we might encounter FX nodes that return zero tensor outputs.
|
||||
if (
|
||||
len(old_res_storage) == 1
|
||||
and len(node_res_storage) == 1
|
||||
and old_res_storage == node_res_storage
|
||||
):
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
new_flattened_res = pytree.tree_leaves(new.meta["fake_result"])
|
||||
new_res_storage = {
|
||||
StorageWeakRef(x._typed_storage())
|
||||
for x in new_flattened_res
|
||||
if isinstance(x, FakeTensor)
|
||||
}
|
||||
if len(new_res_storage) != 1:
|
||||
raise AssertionError(
|
||||
f"Expected 1 storage, got {len(new_res_storage)}"
|
||||
)
|
||||
(new_ref,) = new_res_storage
|
||||
(node_ref,) = node_res_storage
|
||||
# Technically, "old_ref" and all its aliases will remain
|
||||
# in our mapping.
|
||||
# That should be fine though, since we deleted "old"
|
||||
# from the graph at this point.
|
||||
storage_to_nodes[node_ref].update(storage_to_nodes[new_ref])
|
||||
storage_to_nodes[new_ref].update(storage_to_nodes[node_ref])
|
||||
|
||||
# Step 4: delete any _scatter nodes that we de-functionalized
|
||||
# Need to take care not to delete any of these nodes until after *all* modifications
|
||||
# to the graph are finished.
|
||||
for to_delete in all_later_view_inverse_nodes_to_delete:
|
||||
gm.graph.erase_node(to_delete)
|
||||
|
||||
gm.recompile()
|
||||
return gm
|
||||
Reference in New Issue
Block a user