code: def reinplace(gm, *sample_args):
      """
      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 re-used
      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 we 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 re-using 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.
              It's only a problem if "a" (or that view) is later passed
          (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, 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 toe 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:
              b = foo(a)
            After:
              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"
      }
  
      # 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 idx, node in enumerate(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)
                  != 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_name = self_arg.name
              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 re-used 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,
                      )
                      copy_node = 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:
                      new_args = []
                      args = node_to_update.args
  
                      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
                      ):
                          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)
                          }
                          assert len(new_res_storage) == 1
                          (old_ref,) = old_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
  
imports: import torch
  from torch.fx import Node
  from torch.fx._compatibility import compatibility
  from torch._subclasses.fake_tensor import FakeTensorMode, FakeTensor
  from torch.utils._pytree import tree_map_only
  from torch.utils import _pytree
  from torch.multiprocessing.reductions import StorageWeakRef
  import _operator
  from enum import Enum
  import itertools
  from typing import Set, Dict
  from collections import defaultdict
label_desc: Use image node_id 7 for calling a global function with example usage: reinplace(gm) and returns: gm, a, new