# mypy: allow-untyped-defs
import itertools
import logging
import operator
from collections import defaultdict
from collections.abc import Callable, Sequence
from contextlib import nullcontext
from dataclasses import dataclass
from typing import Any, cast

import torch
import torch.fx.node
from torch._C._dynamo.guards import compute_overlapping_tensors
from torch._dispatch.python import enable_python_dispatcher
from torch._dynamo.utils import ReinplaceCounters, ReInplaceTrigger
from torch._guards import detect_fake_mode
from torch._higher_order_ops.triton_kernel_wrap import (
    kernel_side_table,
    triton_kernel_wrapper_functional,
)
from torch._inductor import config, inductor_prims
from torch._inductor.fx_utils import get_node_storage, is_node_realized
from torch._inductor.lowering import (
    inplaceable_foreach_ops as inplaceable_foreach_ops_lowerings,
)
from torch._inductor.virtualized import V
from torch.fx.experimental.symbolic_shapes import (
    compute_unbacked_bindings,
    GuardOnDataDependentSymNode,
    statically_known_true,
    sym_eq,
)
from torch.fx.immutable_collections import immutable_dict, immutable_list
from torch.fx.passes.reinplace import _is_view_op
from torch.utils import _pytree as pytree
from torch.utils._ordered_set import OrderedSet


log = logging.getLogger(__name__)
aten = torch.ops.aten


@dataclass(frozen=True)
class InplaceableOp:
    inplace_op: Callable[..., Any]
    mutated_arg: int | tuple[int, ...]  # Single index or tuple of indices
    extra_check: Callable[[torch.fx.Node], bool] = lambda node: True

    @property
    def mutated_args(self) -> tuple[int, ...]:
        """Return mutated_arg as a tuple for uniform handling."""
        if isinstance(self.mutated_arg, int):
            return (self.mutated_arg,)
        return self.mutated_arg


_SCATTER_OP_TO_VIEW = {
    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,
}
_VIEW_OP_TO_SCATTER = {v: k for k, v in _SCATTER_OP_TO_VIEW.items()}


def graph_call_function(graph: torch.fx.Graph, fn, *args, **kwargs):
    fake_args, fake_kwargs = pytree.tree_map(
        lambda node: node.meta["val"] if isinstance(node, torch.fx.Node) else node,
        (args, kwargs),
    )
    with V.fake_mode:
        fake_result = fn(*fake_args, **fake_kwargs)

    node = graph.call_function(fn, args, kwargs)

    node.meta["val"] = fake_result

    return node


@dataclass
class ViewOp:
    target: torch._ops.OpOverload
    args: tuple[Any, ...]
    kwargs: dict[str, Any]


def _inplace_generalized_scatter(
    inp: torch.Tensor, src: torch.Tensor, view_ops: list[ViewOp]
) -> torch.Tensor:
    tmp = inp
    for view in view_ops:
        fake_args, fake_kwargs = pytree.tree_map(
            lambda node: node.meta["val"] if isinstance(node, torch.fx.Node) else node,
            (view.args, view.kwargs),
        )
        # slice and select can allocate new unbacked symints, but those won't be reflected
        # in the output of this function, hence shall be ignored.
        fake_mode = detect_fake_mode(fake_args)
        with (
            fake_mode.shape_env.ignore_fresh_unbacked_symbols()
            if fake_mode and fake_mode.shape_env
            else nullcontext()
        ):
            tmp = view.target(tmp, *fake_args, **fake_kwargs)
    try:
        tmp.copy_(src)
    except RuntimeError as e:
        raise RuntimeError(
            f"shape error in scatter op, can not broadcast {src.shape} to {tmp.shape}"
        ) from e
    return inp


def _generalized_scatter(
    inp: torch.Tensor, src: torch.Tensor, view_ops: list[ViewOp]
) -> torch.Tensor:
    out = inp.clone()
    return _inplace_generalized_scatter(out, src, view_ops)


def _decompose_scatter_functional_helper(
    graph: torch.fx.Graph,
    inp: torch.Tensor,
    src: torch.Tensor,
    view_ops: list[ViewOp],
) -> torch.fx.Node:
    view_op, view_ops_tail = view_ops[0], view_ops[1:]

    if view_ops_tail:
        view = graph_call_function(
            graph, view_op.target, inp, *view_op.args, **view_op.kwargs
        )
        src = _decompose_scatter_functional_helper(graph, view, src, view_ops[1:])  # type: ignore[assignment]

    return graph_call_function(
        graph,
        _VIEW_OP_TO_SCATTER[view_op.target],
        inp,
        src,
        *view_op.args,
        **view_op.kwargs,
    )


def _decompose_scatter_functional(
    graph: torch.fx.Graph, node: torch.fx.Node
) -> torch.fx.Node:
    """Decompose _generalized_scatter to a sequence of view_scatter operations

    e.g. _generalized_scatter(inp, src, [(aten.slice, 0, 0, 10), (aten.slice, 1, 10, -10)])

    will become

    view = aten.slice(inp, 0, 0, 10)
    view_updated = aten.slice_scatter(view, src, 1, 10, -10)
    inp_updated = aten.slice_scatter(inp, view_updated, 0, 0, 10)
    """
    assert node.target is _generalized_scatter
    return _decompose_scatter_functional_helper(graph, *node.args)  # type: ignore[arg-type]


def _decompose_scatter_mutating(
    graph: torch.fx.Graph, node: torch.fx.Node
) -> torch.fx.Node:
    """Decompose _generalized_scatter using mutations

    e.g. _generalized_scatter(inp, src, [(aten.slice, 0, 0, 10), (aten.slice, 1, 10, -10)])

    will become

    inp_updated = aten.clone(inp)
    slice1 = aten.slice(inp_updated, 0, 0, 10)
    slice2 = aten.slice(slice1, 1, 10, -10)
    slice2.copy_(src)

    """
    assert node.target in (_generalized_scatter, _inplace_generalized_scatter)
    inp, src, view_ops = node.args
    assert not node.kwargs

    if node.target is _generalized_scatter:
        inp = graph_call_function(graph, aten.clone, inp)

    tmp = inp
    for view in view_ops:  # type: ignore[union-attr]
        tmp = graph_call_function(graph, view.target, tmp, *view.args, **view.kwargs)  # type: ignore[union-attr]
        # we need to set unbacked bindings that could have been created in the view ops.
        if (V.fake_mode.shape_env) and (
            symbol_to_path := compute_unbacked_bindings(
                V.fake_mode.shape_env, tmp.meta["val"]
            )
        ):
            tmp.meta["unbacked_bindings"] = symbol_to_path

    graph_call_function(graph, aten.copy_.default, tmp, src)
    return inp  # type: ignore[return-value]


# View ops whose view_scatter op is lowered into mutations anyway,
# so is never a pessimisation to decompose.
_ALWAYS_MUTATING_SCATTER_OPS = OrderedSet(
    [
        aten.as_strided.default,
        aten.diagonal.default,
    ]
)


def scatter_always_uses_mutation(node: torch.fx.Node) -> bool:
    _, _, view_ops = node.args
    view_ops = cast(Sequence[torch.fx.node.Argument], view_ops)
    return any(
        target in _ALWAYS_MUTATING_SCATTER_OPS
        for view in view_ops
        if isinstance(target := getattr(view, "target", None), torch._ops.OpOverload)
    )


def should_reinplace_scatter(node: torch.fx.Node) -> bool:
    """Choose between mutating and functional scatter decompositions

    Reinplacing view scatter ops can be pessimising as it blocks fusion with the
    input or output tensor computations. However, it is still profitable if the
    input and output would have been realized anyway.

    """
    inp, _src, _view_ops = node.args

    # Mutating scatter ops unconditionally realize input and output
    if scatter_always_uses_mutation(node):
        return True

    if is_node_realized(inp) and is_node_realized(node):  # type: ignore[arg-type]
        return True

    # If the output is copied back into the input, this forces both to be
    # realized as the output is a user of the input
    if inp.op in ("placeholder", "get_attr") and any(  # type: ignore[union-attr]
        user.target is aten.copy_.default and user.args[0] is inp for user in node.users
    ):
        return True

    # Otherwise, assume fusions will make functional variants profitable
    return False


def decompose_generalized_scatter(graph: torch.fx.Graph) -> None:
    """Replace _generalized_scatter with normal aten ops"""
    for node in itertools.chain(
        graph.find_nodes(op="call_function", target=_generalized_scatter),
        graph.find_nodes(op="call_function", target=_inplace_generalized_scatter),
    ):
        use_mutation = (
            node.target is _inplace_generalized_scatter
            or scatter_always_uses_mutation(node)
        )

        with graph.inserting_before(node):
            if use_mutation:
                new_node = _decompose_scatter_mutating(graph, node)
            else:
                new_node = _decompose_scatter_functional(graph, node)

        node.replace_all_uses_with(new_node)
        graph.erase_node(node)


def canonicalize_view_scatter_ops(graph: torch.fx.Graph) -> None:
    """
    This canonicalizes view scatter ops into a generalized form, defined as:
      def scatter(inp, src, views):
        tmp = inp.clone()
        for view in views:
          tmp = view(tmp)
        tmp.copy_(src)

    We also fuse consecutive view scatter ops of the form
        a = scatter(view2(self), src, [view1])
        b = scatter(self, a, [view2])
    which can be rewritten as
        b = scatter(self, src, [view2, view1])
        a = view2(b)

    This is both more efficient as we only do a single scatter, and also
    easier to reinplace since there is only one use of `self`
    """

    node_to_view_base: dict[torch.fx.Node, torch.fx.Node] = {}
    node_to_view_op: dict[torch.fx.Node, list[ViewOp]] = defaultdict(list)

    def handle_views(node: torch.fx.Node):
        inp = node.args[0]
        node_to_view_base[node] = node_to_view_base.get(inp, inp)  # type: ignore[arg-type, assignment]
        node_to_view_op[node] = [
            *node_to_view_op[inp],  # type: ignore[index]
            ViewOp(
                node.target,  # type: ignore[arg-type]
                args=node.args[1:],
                kwargs=node.kwargs,
            ),
        ]

    def handle_view_scatter(node: torch.fx.Node):
        assert len(node.args) >= 2
        inp, src = node.args[:2]

        assert isinstance(node.target, torch._ops.OpOverload)
        scatter_view_op = ViewOp(
            _SCATTER_OP_TO_VIEW[node.target],
            args=node.args[2:],
            kwargs=node.kwargs,
        )

        def can_fuse():
            if src.target is not _generalized_scatter:  # type: ignore[union-attr]
                return False
            src_inp, _src_src, _src_scatter_view_op = src.args  # type: ignore[union-attr]

            inp_base = node_to_view_base.get(inp, inp)  # type: ignore[arg-type]
            src_base = node_to_view_base.get(src_inp, src_inp)  # type: ignore[arg-type]
            return inp_base is src_base and node_to_view_op[src_inp] == [  # type: ignore[index]
                *node_to_view_op[inp],  # type: ignore[index]
                scatter_view_op,
            ]

        if not can_fuse():
            with graph.inserting_before(node):
                new_node = graph_call_function(
                    graph,
                    _generalized_scatter,
                    inp,
                    src,
                    [scatter_view_op],
                )
            node.replace_all_uses_with(new_node)
            graph.erase_node(node)
            return

        _src_inp, src_src, src_scatter_view_op = src.args  # type: ignore[union-attr]
        with graph.inserting_before(src):  # type: ignore[arg-type]
            new_node = graph_call_function(
                graph,
                _generalized_scatter,
                inp,
                src_src,
                [scatter_view_op, *src_scatter_view_op],  # type: ignore[misc]
            )
            node.replace_all_uses_with(new_node)
            graph.erase_node(node)

            if src.users:  # type: ignore[union-attr]
                new_src = graph_call_function(
                    graph,
                    _SCATTER_OP_TO_VIEW[node.target],
                    new_node,
                    *node.args[2:],
                    **node.kwargs,
                )

                handle_views(new_src)
                src.replace_all_uses_with(new_src)  # type: ignore[union-attr]

            graph.erase_node(src)  # type: ignore[arg-type]

    for node in graph.nodes:
        if _is_view_op(node.target):
            handle_views(node)
        elif node.target in _SCATTER_OP_TO_VIEW:
            handle_view_scatter(node)


inplaceable_ops: dict[Callable[..., Any], InplaceableOp] = {
    aten.index_put.default: InplaceableOp(aten.index_put_.default, 0),
    aten._unsafe_index_put.default: InplaceableOp(inductor_prims._unsafe_index_put_, 0),
    _generalized_scatter: InplaceableOp(
        _inplace_generalized_scatter,
        0,
        extra_check=should_reinplace_scatter,
    ),
}

try:
    c10d_functional = torch.ops._c10d_functional
    inplaceable_collective_ops: dict[Callable[..., Any], InplaceableOp] = {
        c10d_functional.all_reduce.default: InplaceableOp(
            c10d_functional.all_reduce_.default, 0
        ),
        c10d_functional.all_reduce_coalesced.default: InplaceableOp(
            c10d_functional.all_reduce_coalesced_.default, 0
        ),
    }
    inplaceable_ops.update(inplaceable_collective_ops)
except AttributeError:
    # _c10d_functional ops are only available when torch
    # is built with USE_DISTRIBUTED=1.
    pass

inplaceable_foreach_ops: dict[torch._ops.OpOverload, InplaceableOp] = {}
for outplace_op, inplace_op in inplaceable_foreach_ops_lowerings.items():
    inplaceable_foreach_ops[outplace_op] = InplaceableOp(inplace_op, 0)


inplaceable_triton_ops = OrderedSet([triton_kernel_wrapper_functional])


# Operators that don't depend on the tensor data
META_ONLY_OPS = OrderedSet(
    [
        aten.sym_size.int,
        aten.sym_stride.int,
        aten.sym_numel.default,
        aten.sym_storage_offset.default,
    ]
)


def _get_view_base(node: torch.fx.Node) -> torch.fx.Node:
    while _is_view_op(node.target) and node.args:
        base = node.args[0]
        if not isinstance(base, torch.fx.Node):
            break
        node = base
    return node


def _same_tensor_metadata(lhs: torch.fx.Node, rhs: torch.fx.Node) -> bool:
    lhs_val = lhs.meta.get("val")
    rhs_val = rhs.meta.get("val")
    if not isinstance(lhs_val, torch.Tensor) or not isinstance(rhs_val, torch.Tensor):
        return False

    def same_value(lhs_value, rhs_value) -> bool:
        return statically_known_true(sym_eq(lhs_value, rhs_value))

    def same_sequence(lhs_values, rhs_values) -> bool:
        return len(lhs_values) == len(rhs_values) and all(
            same_value(lhs_value, rhs_value)
            for lhs_value, rhs_value in zip(lhs_values, rhs_values)
        )

    return (
        same_sequence(lhs_val.size(), rhs_val.size())
        and same_sequence(lhs_val.stride(), rhs_val.stride())
        and same_value(lhs_val.storage_offset(), rhs_val.storage_offset())
    )


def _is_layout_preserving_view_copy_back(
    dst: torch.fx.Node,
    src: torch.fx.Node,
    mutated_arg: torch.fx.Node,
    src_base: torch.fx.Node,
) -> bool:
    return _same_tensor_metadata(src_base, mutated_arg) and _same_tensor_metadata(
        src, dst
    )


def reinplace_inplaceable_ops_core(graph: torch.fx.Graph) -> None:
    """
    Reinplaces in-placeable operations.
    If there are no uses of a view of the mutated arg after the current node,
    it is possible to inplace the op.
    This above algorithm could be justified by observing side effects. While
    we traverse the graph in forwards direction, only latter nodes could view
    side effects of the current node. If the current node is not used later as
    well as no view of this node is used later in the graph, then it is safe to
    inplace as there would be no way to observe the side effects.
    This condition is slightly different for graph inputs where they can only
    be inplaced if the above condition is true and there's a copy_ in the
    epilogue that signals that the caller wants to observe the mutation.

    Unlike JIT Inductor, AOTInductor currently unlifts weights and buffers from
    input args, so instead of checking mutation on placeholder, AOTInductor
    checks mutation on get_attr. This is subject to change in future.
    """

    copy_args_to_copy_nodes = {}
    # maps argument to the first copy_ node that mutates it.
    copy_nodes = {}
    # maps (view_arg, inplaceable_op_node) to copy_ nodes that copy a view of
    # the op result back into the graph input that view_arg aliases.
    copy_args_to_copy_nodes_via_views = {}
    mutated_inputs = OrderedSet[Any]()
    storage_to_nodes = defaultdict(list)
    node_order: dict[Any, int] = {}
    for i, node in enumerate(reversed(graph.nodes)):
        node_order[node] = len(graph.nodes) - i - 1
        storage_to_nodes[get_node_storage(node)].append(node)
        if node.target is aten.copy_.default and node.args[0].op in (
            "placeholder",
            "get_attr",
        ):
            dst = node.args[0]
            src = node.args[1]
            # If the target is a getitem and it indexes a possible clone,
            # then skip over it
            if src.target is operator.getitem and (
                (
                    src.args[0].target == triton_kernel_wrapper_functional
                    and src.args[0].kwargs["kwargs"][src.args[1]] == node.args[0]
                )
                or (src.args[0].target in inplaceable_foreach_ops)
                or (src.args[0].target is torch.ops.higher_order.auto_functionalized)
            ):
                src = src.args[0]

            copy_args_to_copy_nodes[(dst, src)] = node
            copy_nodes[dst] = node

            mutated_inputs.add(node.args[0])

            src_base = _get_view_base(src)
            if src_base is not src and isinstance(src_base.target, Callable):
                inplaceable_op = inplaceable_ops.get(src_base.target)
                if inplaceable_op is not None:
                    for mutated_arg_idx in inplaceable_op.mutated_args:
                        mutated_arg = src_base.args[mutated_arg_idx]
                        if (
                            isinstance(mutated_arg, torch.fx.Node)
                            and _get_view_base(mutated_arg) is dst
                            and _is_layout_preserving_view_copy_back(
                                dst, src, mutated_arg, src_base
                            )
                        ):
                            copy_args_to_copy_nodes_via_views[
                                (mutated_arg, src_base)
                            ] = node

    def any_use_of_views_after_node(node, shared_view_nodes, *, copy_node, mutated_arg):
        node_loc = node_order[node]
        copy_node_loc = node_order[copy_node] if copy_node is not None else None

        def is_meta_only_user(node):
            if _is_view_op(node.target):
                return all(is_meta_only_user(u) for u in node.users)
            return node.target in META_ONLY_OPS

        for view in shared_view_nodes:
            for user in view.users:
                user_loc = node_order[user]
                # Skip all users before node
                if user_loc <= node_loc:
                    continue
                # Ignore uses after the copy_ epilogue node, where the input
                # has already been mutated anyway
                if copy_node_loc is not None and copy_node_loc <= user_loc:
                    continue
                # Reinplacing does not change shape metadata
                if is_meta_only_user(user):
                    continue
                # If our graph looks like:
                # foo(mutated_arg)
                # mutated_arg.copy_(other)
                # then it's safe for us to reinplace foo because mutated_arg
                # will get overwritten anyways.
                if (
                    user.target is torch.ops.aten.copy_.default
                    and mutated_arg is user.args[0]
                ):
                    continue
                return True
        return False

    def can_inplace(node, mutated_arg):
        # ls should be a list of tensors that all shares the same storage.
        def _overlap(ls) -> bool:
            try:
                return len(compute_overlapping_tensors(ls)) != 0
            except GuardOnDataDependentSymNode:
                # If we fail with data dependent error we assume they all overlap.
                return True

        if isinstance(mutated_arg, (list, tuple)):
            # TODO Using _overlap here causes a several issues.
            unique_storages = OrderedSet(get_node_storage(arg) for arg in mutated_arg)
            if len(unique_storages) != len(mutated_arg):
                # At least two Tensors in mutated_arg alias each other, so we can't reinplace it.
                # We can probably do better (that is, reinplace one of them and clone the other)
                # but that requires more work and mutable List[Tensor] are not that common.
                return False
            return all(can_inplace(node, arg) for arg in mutated_arg)

        if get_node_storage(mutated_arg) is None:
            return False

        shared_view_nodes = storage_to_nodes[get_node_storage(mutated_arg)]

        # Only keep tensor that might overlap with mutated_arg.
        shared_view_nodes = [
            v
            for v in shared_view_nodes
            if _overlap([mutated_arg.meta["val"], v.meta["val"]])
        ]

        if mutated_arg.op in ("placeholder", "get_attr"):
            # Get the first copy_ node that mutates the mutated_arg.
            copy_node = copy_nodes.get(mutated_arg)
            if copy_node is None:
                # There is no copy_ back to the candidate mutated_arg (which is a graph input).
                # Therefore the semantics of the program are that it does not mutate
                # mutated_arg, so we cannot re-inplace it.
                return False
            if any_use_of_views_after_node(
                node, shared_view_nodes, copy_node=copy_node, mutated_arg=mutated_arg
            ):
                return False

            return True
        elif any(view.op in ("placeholder", "get_attr") for view in shared_view_nodes):
            # If mutated_arg is a view of a graph input, we can only reinplace
            # when a later copy_ writes this op's result back to that input.
            copy_node = copy_args_to_copy_nodes_via_views.get((mutated_arg, node))
            if copy_node is None:
                return False

            mutated_arg_base = copy_node.args[0]
            copy_src = copy_node.args[1]
            if not (
                isinstance(mutated_arg_base, torch.fx.Node)
                and isinstance(copy_src, torch.fx.Node)
                and _is_layout_preserving_view_copy_back(
                    mutated_arg_base, copy_src, mutated_arg, node
                )
            ):
                return False

            if any_use_of_views_after_node(
                node,
                shared_view_nodes,
                copy_node=copy_node,
                mutated_arg=mutated_arg_base,
            ):
                return False
            return True
        else:
            return not any_use_of_views_after_node(
                node, shared_view_nodes, copy_node=None, mutated_arg=mutated_arg
            )

    def copy_node_for_reinplaced_arg(node, mutated_arg):
        copy_node = copy_args_to_copy_nodes.get((mutated_arg, node))
        if copy_node is None:
            copy_node = copy_args_to_copy_nodes_via_views.get((mutated_arg, node))
            if copy_node is not None:
                copy_dst = copy_node.args[0]
                copy_src = copy_node.args[1]
                if not (
                    isinstance(copy_dst, torch.fx.Node)
                    and isinstance(copy_src, torch.fx.Node)
                    and _is_layout_preserving_view_copy_back(
                        copy_dst, copy_src, mutated_arg, node
                    )
                ):
                    return None
        return copy_node

    def all_can_inplace(node, mutated_args):
        return all(can_inplace(node, arg) for arg in mutated_args)

    def log_inplace_results(
        old_node_name,
        node_name,
        old_tensors_to_clone,
        tensors_to_clone,
        missed_args,
        missed_nodes,
        trigger,
    ):
        # Total size of possibly_missed_reinplacing_opportunities for tensors with static shapes.
        missed_bytes = 0

        def bytes(node):
            t = node.meta.get("val", None)
            if (
                t is not None
                and isinstance(t.element_size(), int)
                and isinstance(t.numel(), int)
            ):
                return t.element_size() * t.numel()
            else:
                return 0

        for node in missed_nodes:
            if isinstance(node, (list, tuple)):
                for n in node:
                    missed_bytes += bytes(n)
            else:
                missed_bytes += bytes(node)

        log.info(
            "Reinplace for %s (wrapped by %s): candidates=%s, successfully reinplaced=%s, "
            "remain clones=%s, missed opportunities=%s, missed static bytes=%s.",
            old_node_name,
            node_name,
            old_tensors_to_clone,
            [
                tensor_idx
                for tensor_idx in old_tensors_to_clone
                if tensor_idx not in tensors_to_clone
            ],
            tensors_to_clone,
            missed_args,
            missed_bytes,
        )

        ReinplaceCounters.add_missed_opportunities(trigger, len(missed_args))
        ReinplaceCounters.add_missed_bytes(trigger, missed_bytes)

    replace_dict: dict[torch.fx.Node, torch.fx.Node] = {}

    def reinplace_and_refine_tensors_to_clone(
        old_tensors_to_clone, kwargs, node_name, old_node_name, trigger
    ):
        tensors_to_clone: list[str] = []
        storage_of_reinplaced_args = OrderedSet[int | None]()

        # Those used to count possibly_missed_reinplacing_opportunities
        missed_nodes = []
        missed_args = []

        # TODO this logic can be made more precise using _overlap
        def tensor_with_same_storage_already_reinplaced(arg):
            if isinstance(arg, (list, tuple)):
                return any(
                    get_node_storage(a) in storage_of_reinplaced_args for a in arg
                )
            return get_node_storage(mutated_arg) in storage_of_reinplaced_args

        for arg in old_tensors_to_clone:
            assert arg in kwargs

            mutated_arg = kwargs[arg]

            # Let's say we have:
            # - op(x, y) that mutates both x and y
            # - new_x, new_y = functional_op(x, y) is the functional variant
            # If we are presented with functional_op(x, x), we must not reinplace
            # this into op(x, x), because then it would be writing to the same Tensor.
            # Instead, it's OK to reinplace one of them and to clone the other:
            # >>> y = x.clone()
            # >>> op(x, y)
            # This also applies if we have views: functional_op(x, x[0])
            # should not reinplace into op(x, x[0]).
            should_attempt_reinplace = not tensor_with_same_storage_already_reinplaced(
                mutated_arg
            )
            if should_attempt_reinplace and can_inplace(node, mutated_arg):
                # In general, we probably do not need those optimizations.
                copy_node = copy_node_for_reinplaced_arg(node, mutated_arg)
                if copy_node is not None:
                    replace_dict[copy_node] = copy_node.args[0]
                if trigger != ReInplaceTrigger.AUTO_FUNC_V2:
                    for user in node.users:
                        # For auto_functionalize_v2, arg is the index of the base, where base at index i corresponds to
                        # output atindex size(out)+i.
                        # This used to compare string with integers before for auto_functionalize_v2. Not sure
                        # if it was needed for inplaceable_triton_ops?
                        if user.target is operator.getitem and user.args[1] == arg:
                            replace_dict[user] = mutated_arg

                if isinstance(mutated_arg, (list, tuple)):
                    for a in mutated_arg:
                        storage_of_reinplaced_args.add(get_node_storage(a))
                else:
                    storage_of_reinplaced_args.add(get_node_storage(mutated_arg))
            else:
                if should_attempt_reinplace:
                    missed_args.append(arg)
                    missed_nodes.append(mutated_arg)

                tensors_to_clone.append(arg)

        log_inplace_results(
            old_node_name,
            node_name,
            old_tensors_to_clone,
            tensors_to_clone,
            missed_args,
            missed_nodes,
            trigger,
        )
        return tensors_to_clone

    for node in graph.nodes:
        if (inplaceable_op := inplaceable_ops.get(node.target)) is not None:
            # Check if ALL mutated args can be inplaced
            # Only convert if we don't need to clone any tensor
            mutated_args = [node.args[idx] for idx in inplaceable_op.mutated_args]
            if all_can_inplace(node, mutated_args) and inplaceable_op.extra_check(node):
                for mutated_arg in mutated_args:
                    copy_node = copy_node_for_reinplaced_arg(node, mutated_arg)
                    if copy_node is not None:
                        replace_dict[copy_node] = copy_node.args[0]
                node.target = inplaceable_op.inplace_op
        elif node.target is torch.ops.higher_order.auto_functionalized_v2:
            _mutable_op = node.args[0]
            kwargs = node.kwargs

            if isinstance(
                _mutable_op, torch._ops.OpOverload
            ) and torch._library.utils.is_out(_mutable_op):
                # Out args are write-only, always safe to reinplace (no clones needed)
                node.meta["only_clone_these_tensors"] = []
            else:
                all_bases = kwargs["_all_bases"]
                bases_to_clone = range(len(all_bases))
                base_tensors_dct = dict(enumerate(all_bases))
                new_bases_to_clone: list[int] = reinplace_and_refine_tensors_to_clone(
                    bases_to_clone,
                    base_tensors_dct,
                    node.name,
                    _mutable_op._name,
                    ReInplaceTrigger.AUTO_FUNC_V2,
                )
                # Stash the metadata. There is a pass later on where we decompose
                # auto_functionalized into clones + a mutable op; this metadata
                # tells the decomp to only clone the following inputs
                node.meta["only_clone_these_tensors"] = new_bases_to_clone
        elif node.target is torch.ops.higher_order.auto_functionalized:
            _mutable_op = node.args[0]
            from torch._higher_order_ops.auto_functionalize import get_mutable_args

            tensors_to_clone, _ = get_mutable_args(_mutable_op)
            # Don't try to reinplace Tensor | None args that are None.
            tensors_to_clone = [
                t for t in tensors_to_clone if node.kwargs[t] is not None
            ]
            tensors_to_clone = reinplace_and_refine_tensors_to_clone(
                tensors_to_clone,
                node.kwargs,
                node.name,
                _mutable_op._name,
                ReInplaceTrigger.AUTO_FUNC_V1,
            )

            # Stash the metadata. There is a pass later on where we decompose
            # auto_functionalized into clones + a mutable op; this metadata
            # tells the decomp to only clone the following inputs
            node.meta["only_clone_these_tensors"] = tensors_to_clone
        elif node.target is torch.ops.higher_order.with_effects:
            # Handle effectful ops wrapped with with_effects
            # args[0] is the token, args[1] is the inner op, args[2:] are the op's args
            inner_op = node.args[1]
            log.debug(
                "reinplace: checking with_effects node with inner_op=%s", inner_op
            )
            if (inplaceable_op := inplaceable_ops.get(inner_op)) is not None:
                log.debug("reinplace: found inplaceable_op for %s", inner_op)
                # Get the mutated arg indices, offset by 2 (token + op)
                mutated_arg_indices = inplaceable_op.mutated_args

                # Build flat list of tensors for can_inplace check
                # and a mapping of output index -> replacement tensor(s)
                mutated_tensors_flat = []
                output_idx_to_replacement: dict[int, Any] = {}

                for position, idx in enumerate(mutated_arg_indices):
                    actual_idx = idx + 2  # offset for token and op
                    assert actual_idx < len(node.args), (
                        f"mutated arg idx {actual_idx} out of range {len(node.args)}"
                    )
                    arg = node.args[actual_idx]

                    # Output index is position + 1 (index 0 is the token)
                    output_idx = position + 1
                    output_idx_to_replacement[output_idx] = arg

                    # Flatten for can_inplace check
                    if isinstance(arg, (list, tuple)):
                        mutated_tensors_flat.extend(arg)
                    else:
                        mutated_tensors_flat.append(arg)

                # Check if all mutated args can be inplaced
                can_inplace_all = all_can_inplace(node, mutated_tensors_flat)

                log.debug(
                    "reinplace with_effects: mutated_tensors=%s, can_inplace_all=%s",
                    [str(a) for a in mutated_tensors_flat],
                    can_inplace_all,
                )

                if can_inplace_all and inplaceable_op.extra_check(node):
                    log.debug(
                        "reinplace with_effects: converting %s -> %s",
                        inner_op,
                        inplaceable_op.inplace_op,
                    )
                    # Update the inner op to inplace version
                    node.update_arg(1, inplaceable_op.inplace_op)

                    # The output structure changes: functional returns (token, tensors),
                    # inplace returns (token, None). We need to redirect tensor uses
                    # to the input tensors.

                    def get_index_from_node(n):
                        """Extract the index from a getitem node."""
                        if n.target is operator.getitem:
                            return n.args[1]
                        return None

                    def is_getitem_node(n, parent):
                        """Check if node n is a getitem indexing into parent."""
                        return n.target is operator.getitem and n.args[0] is parent

                    def replace_and_collect(current_node, replacement_tensors):
                        """
                        Collect replacements for getitem nodes into replace_dict.
                        Nodes are added in child-first order so children are erased before parents.
                        """
                        # Find all users that are getitem nodes indexing into current_node
                        getitem_users = [
                            u
                            for u in current_node.users
                            if is_getitem_node(u, current_node)
                        ]

                        if not getitem_users:
                            # Leaf node - add to replace_dict with actual replacement
                            if isinstance(replacement_tensors, (list, tuple)):
                                if len(replacement_tensors) == 1:
                                    replace_dict[current_node] = replacement_tensors[0]
                                    return True
                                else:
                                    # Multiple tensors but no indexing - can't replace
                                    return False
                            else:
                                replace_dict[current_node] = replacement_tensors
                                return True

                        # Process children first (so they're added to replace_dict before parent)
                        all_children_replaced = True
                        first_replacement = None
                        for getitem_user in getitem_users:
                            idx = get_index_from_node(getitem_user)
                            if idx is None or not isinstance(idx, int):
                                all_children_replaced = False
                                continue

                            if not isinstance(replacement_tensors, (list, tuple)):
                                all_children_replaced = False
                                continue

                            if idx >= len(replacement_tensors):
                                all_children_replaced = False
                                continue

                            if first_replacement is None:
                                first_replacement = replacement_tensors[idx]

                            if not replace_and_collect(
                                getitem_user, replacement_tensors[idx]
                            ):
                                all_children_replaced = False

                        # Add this node to replace_dict after children (even if it has non-getitem users)
                        # Non-getitem users will have their input replaced via replace_all_uses_with
                        if all_children_replaced and first_replacement is not None:
                            replace_dict[current_node] = first_replacement

                        return all_children_replaced

                    # Find getitem nodes that extract tensor results
                    # Use the output_idx_to_replacement mapping built above
                    for user in list(node.users):
                        if not is_getitem_node(user, node):
                            continue
                        idx = get_index_from_node(user)
                        if idx is None or idx not in output_idx_to_replacement:
                            continue
                        replacement = output_idx_to_replacement[idx]
                        replace_and_collect(user, replacement)
        elif node.target in inplaceable_triton_ops:
            kernel_idx = node.kwargs["kernel_idx"]
            kernel = kernel_side_table.get_kernel(kernel_idx)
            from triton.runtime.autotuner import Autotuner
            from triton.runtime.jit import JITFunction

            if isinstance(kernel, JITFunction):
                kernel_name = kernel.fn.__name__
            elif isinstance(kernel, Autotuner):
                if config.is_fbcode():
                    # Autotuner has different implementations for AMD and NV
                    if torch.version.hip is None:
                        kernel_name = kernel.base_fn.__name__
                    else:
                        kernel_name = kernel.fn.__name__
                else:
                    kernel_name = kernel.base_fn.__name__
            else:
                raise AssertionError("Unknown triton kernel type")

            # inplaceable_triton_ops take an additional argument called
            # tensors_to_clone which contain a list of tensors to clone
            # This pass iterates over them and sees which ones are safe
            # to eliminate (i.e. no longer need the clones)
            tensors_to_clone = reinplace_and_refine_tensors_to_clone(
                node.kwargs["tensors_to_clone"],
                node.kwargs["kwargs"],
                node.name,
                kernel_name,
                ReInplaceTrigger.TRITON_OPS,
            )

            kwargs = dict(node.kwargs)
            kwargs["tensors_to_clone"] = tensors_to_clone
            node.kwargs = immutable_dict(kwargs)
            if "eager_input_vals" in node.meta:
                # We changed the kwargs, so we need to update eager_input_vals
                # to something sane.
                args, kwargs = node.meta["eager_input_vals"]
                new_kwargs = {**kwargs}
                new_kwargs["tensors_to_clone"] = immutable_list(tensors_to_clone)
                new_kwargs = immutable_dict(new_kwargs)
                node.meta["eager_input_vals"] = (args, new_kwargs)
        elif (inplaceable_op := inplaceable_foreach_ops.get(node.target)) is not None:
            mutated_args = node.args[inplaceable_op.mutated_arg]

            if not all((arg, node) in copy_args_to_copy_nodes for arg in mutated_args):
                continue

            if can_inplace(node, mutated_args):
                for arg in mutated_args:
                    copy_node = copy_args_to_copy_nodes[(arg, node)]
                    replace_dict[copy_node] = copy_node.args[0]

                node.target = inplaceable_op.inplace_op
    for node, replacement in replace_dict.items():
        while replacement in replace_dict:
            replacement = replace_dict[replacement]
        replace_dict[node] = replacement

        node.replace_all_uses_with(replacement)
        graph.erase_node(node)


def reinplace_inplaceable_ops(
    fake_tensor_updater: torch._inductor.fx_utils.FakeTensorUpdater,
    graph: torch.fx.Graph,
) -> None:
    with enable_python_dispatcher():
        canonicalize_view_scatter_ops(graph)
        # canonicalize_view_scatter_ops adds new operations to the graph.
        # We run fake_tensor_updater to update the alias information.
        # Correct alias information is required for `reinplace_inplaceable_ops_core`.
        fake_tensor_updater.incremental_update()
        reinplace_inplaceable_ops_core(graph)
        decompose_generalized_scatter(graph)
