# mypy: allow-untyped-defs
import contextlib
import operator
import warnings
from collections import defaultdict
from collections.abc import (
    Callable,
    Collection,
    Container,
    Generator,
    Iterable,
    Mapping,
)
from dataclasses import dataclass
from functools import partial
from itertools import chain
from typing import Any

import sympy

import torch
import torch.fx
from torch._dispatch.python import enable_python_dispatcher
from torch._inductor.fx_passes.control_dependencies import control_deps
from torch._subclasses.fake_tensor import FakeTensorMode
from torch.fx.experimental.symbolic_shapes import (
    compute_unbacked_bindings,
    rebind_unbacked,
    statically_known_true,
    sym_eq,
)
from torch.utils import _pytree as pytree
from torch.utils._ordered_set import OrderedSet
from torch.utils._pytree import tree_map
from torch.utils.flop_counter import flop_registry

from .virtualized import V


# Check the pattern: (nn.module, F.function/torch.Tensor.method) matched.
# Works for length 2 patterns with 1 module and 1 function/method.
def matches_module_function_pattern(
    pattern: tuple[type[torch.nn.modules.Module], Callable[..., Any]],
    node: torch.fx.Node,
    modules: dict[str, torch.nn.modules.Module],
) -> bool:
    if len(node.args) == 0:
        return False
    if not isinstance(node.args[0], torch.fx.Node) or not isinstance(
        node, torch.fx.Node
    ):
        return False
    # the first node is call_module
    if node.args[0].op != "call_module":
        return False
    if not isinstance(node.args[0].target, str):
        return False
    if node.args[0].target not in modules:
        return False
    if type(modules[node.args[0].target]) is not pattern[0]:
        return False
    # the second node is call_function or call_method
    if node.op != "call_function" and node.op != "call_method":
        return False
    if node.target != pattern[1]:
        return False
    # make sure node.args[0] output is only used by current node.
    if len(node.args[0].users) > 1:
        return False
    return True


def _is_fake_tensor_same(
    new: Any,
    old: Any,
    existing_storages: Mapping[int, int],
    *,
    check_strides: bool = True,
    check_storage: bool = True,
    node: torch.fx.Node | None = None,
    recursive_ids: OrderedSet[tuple[int, int]] | None = None,
) -> bool:
    """Validate that two FakeTensors (or iterables thereof) are the same, including
    storage locations if enabled.

    check_strides: disabling this flag will remove checks for striding.
    check_storage: disabling this flag will remove checks for storage offset and
    location.  This is useful for subgraph argument and output updating, where storage
    location can change without invalidating the subgraph.
    recursive_ids: This is only for use when recursing through collections, and must not
    be supplied by users of this function."""

    def is_intlist_same(new, old):
        return statically_known_true(sym_eq(new, old))

    if type(new) is not type(old):
        return False

    if not isinstance(new, torch.Tensor):
        if new is None:
            return old is None

        if isinstance(new, Collection):
            if recursive_ids is None:
                recursive_ids = OrderedSet()

            id_pair = (id(new), id(old))
            visited = id_pair in recursive_ids
            recursive_ids.add(id_pair)

            # If we've visited this exact check before while recursing, all elements in
            # this collection have already been validated (or will be validated in the
            # future) by a call at a different layer of recursion.
            return visited or (
                len(new) == len(old)
                and all(
                    _is_fake_tensor_same(
                        new_i,
                        old_i,
                        existing_storages,
                        check_strides=check_strides,
                        check_storage=check_storage,
                        node=node,
                        recursive_ids=recursive_ids,
                    )
                    for new_i, old_i in zip(new, old)
                )
            )

        if isinstance(new, torch.types.py_sym_types):
            return (
                new.node.shape_env._maybe_evaluate_static(
                    sympy.Eq(new.node.expr, old.node.expr)
                )
                == sympy.true
            )

        # If this is not a Collection, a sym-type, or a Tensor, then check for Python
        # equality.  This should be fairly conservative, since an object with no
        # implemented __eq__ method will compare IDs.
        return new == old

    if new.layout != old.layout or not is_intlist_same(new.shape, old.shape):
        return False

    if new.device != old.device:
        return False

    if (
        check_strides
        and new.layout == torch.strided
        and not is_intlist_same(new.stride(), old.stride())
    ):
        return False

    if not check_storage:
        return True

    if not statically_known_true(
        new.storage_offset() == old.storage_offset()
    ) or get_storage(new) != get_storage(old):
        return False

    def any_user_may_alias(node):
        if not isinstance(node.meta["val"], torch.Tensor):
            # analysis too complicated on lists, can support in the future
            return True

        for user in node.users:
            if not (
                isinstance(user.target, torch._ops.OperatorBase)
                or user.target
                is torch._inductor.fx_passes.reinplace._generalized_scatter
            ):
                return True

            if isinstance(user.target, torch._ops.HigherOrderOperator):
                # HOPs that survive until inductor are all non-aliasing HOPs.  We will
                # likely never support HOPs that are aliasing.
                continue

            # Strategy: do a FakeTensor prop, see if the storage aliases.  If Inductor
            # ever gets tighter invariants on OpOverloads (that is, we ban things like
            # torch.ops.aten.reshape calls in the graph), then this could just be a fast
            # schema lookup.
            is_valid, args, kwargs = get_fake_args_kwargs(user)
            if not is_valid:
                return True

            with (
                V.fake_mode,
                enable_python_dispatcher(),
                contextlib.ExitStack() as stack,
            ):
                # Ignore unbacked symbols (if they exist): we're making this FakeTensor
                # and then throwing it away.
                if (shape_env := V.fake_mode.shape_env) is not None:
                    stack.enter_context(shape_env.ignore_fresh_unbacked_symbols())
                new_fake_tensor = user.target(*args, **kwargs)

            if not isinstance(new_fake_tensor, torch.Tensor):
                # analysis too complicated on lists, can support in the future
                return True

            if get_storage(new_fake_tensor) == get_storage(node.meta["val"]):
                return True

        return False

    # This is the case where it returns a completely fresh storage that's used nowhere
    # else.  If the FakeTensor's storage is fresh and none of the node's users can alias
    # it, then we don't need to update this node.
    if (
        existing_storages[get_storage(old)] == 1
        and get_storage(new) not in existing_storages
        and not any_user_may_alias(node)
    ):
        return True

    return False


def _extract_subgraphs_and_args(
    node: torch.fx.Node,
    valid_subgraphs: Container[torch.fx.GraphModule],
    *args: Any,
    **kwargs: Any,
) -> Generator[tuple[torch.fx.GraphModule, tuple[Any, ...] | None]]:
    """HOPs that invoke subgraphs take a number of different forms.  This function
    regularizes them, yielding subgraphs from the args and kwargs coupled with
    appropriate args to update those subgraphs.

    If the second yielded value is None, this function was unable to determine what args
    to pass to the subgraph."""

    def get_subgraph_args(
        subgraph: torch.fx.GraphModule,
    ) -> tuple[torch.Tensor, ...]:
        return tuple(
            get_fake(n, subgraph) for n in subgraph.graph.find_nodes(op="placeholder")
        )

    def is_integer(d: torch.dtype) -> bool:
        return not d.is_floating_point and not d.is_complex

    if node.target is torch.ops.higher_order.associative_scan:
        # Associative scan operates on slices of xs (see: scan), but multiple slices.
        # Use the same slice twice to account for cases where only a single slice is
        # input.
        yield args[0], (*(a[0] for a in args[1]), *(a[0] for a in args[1]), *args[2])
    elif node.target is torch.ops.higher_order.cond:
        subgraph_args = tuple(args[3])
        yield args[1], subgraph_args
        yield args[2], subgraph_args
    elif node.target is torch.ops.higher_order.flex_attention:
        # flex_attention currently yields two subgraphs.  The first (the score
        # calculation) has a signature starting with a scalar argument of the same dtype
        # as the query (args[0]), followed by four scalar integer arguments and optional
        # additional tensors.  The second (the mask calculation) looks similar, with
        # four scalar integer arguments and optional additional (different) tensors.  We
        # assume this format here (adding some defensive assertions), which means that
        # the only inputs we may need to update are the optional additional tensors. See
        # torch._higher_order_ops.flex_attention._math_attention_inner for more details.

        score_subgraph: torch.fx.GraphModule = args[3]
        score_subgraph_args = get_subgraph_args(score_subgraph)
        mask_subgraph = args[4][-1]
        mask_subgraph_args = get_subgraph_args(mask_subgraph)

        integer_arg_dtypes = OrderedSet[torch.dtype](
            a.dtype for a in chain(score_subgraph_args[1:5], mask_subgraph_args[:4])
        )
        assert (
            score_subgraph_args[0].dtype == args[0].dtype
            and len(integer_arg_dtypes) == 1
            and is_integer(integer_arg_dtypes.pop())
            and all(
                len(a.size()) == 0
                for a in chain(score_subgraph_args[:5], mask_subgraph_args[:4])
            )
        ), "flex_attention subgraph arg format has changed!"

        yield score_subgraph, (*score_subgraph_args[:5], *args[7])
        yield args[4][-1], (*mask_subgraph_args[:4], *args[8])
    elif node.target is torch.ops.higher_order.flex_attention_backward:
        fw_subgraph = args[7]
        fw_subgraph_args = get_subgraph_args(fw_subgraph)
        joint_subgraph = args[8]
        joint_subgraph_args = get_subgraph_args(joint_subgraph)
        mask_subgraph = args[9][-1]
        mask_subgraph_args = get_subgraph_args(mask_subgraph)

        integer_arg_dtypes = OrderedSet[torch.dtype](
            a.dtype
            for a in chain(
                fw_subgraph_args[1:5],
                joint_subgraph_args[1:5],
                mask_subgraph_args[:4],
            )
        )
        assert (
            fw_subgraph_args[0].dtype == args[0].dtype
            and joint_subgraph_args[0].dtype == args[0].dtype
            and len(integer_arg_dtypes) == 1
            and is_integer(integer_arg_dtypes.pop())
            and all(
                len(a.size()) == 0
                for a in chain(
                    fw_subgraph_args[:5],
                    joint_subgraph_args[:6],
                    mask_subgraph_args[:4],
                )
            )
        ), "flex_attention_backward subgraph arg format has changed!"

        if fw_subgraph in valid_subgraphs:
            yield fw_subgraph, (*fw_subgraph_args[:5], *args[12])
        if joint_subgraph in valid_subgraphs:
            yield joint_subgraph, (*joint_subgraph_args[:6], *args[12])
        if mask_subgraph in valid_subgraphs:
            yield mask_subgraph, (*mask_subgraph_args[:4], *args[13])
    elif node.target in (
        torch.ops.higher_order.foreach_map,
        torch.ops.higher_order.invoke_quant_packed,
        torch.ops.higher_order.invoke_quant,
    ):
        yield args[0], tuple(args[1:])
    elif node.target is torch.ops.higher_order.invoke_subgraph:
        yield args[0], tuple(args[2:])
    elif node.target is torch.ops.higher_order.map_impl:
        # Map is applied over slices from the first dimension of each value in args[1].
        yield args[0], (*(a[0] for a in args[1]), *args[2])
    elif node.target is torch.ops.higher_order.scan:
        # Scans accept a dim keyword, but the dimensions will be reordered so that at
        # this point we always scan over dim 0.
        yield args[0], (*args[1], *(a[0] for a in args[2]), *args[3])
    elif node.target in (
        torch.ops.higher_order.while_loop,
        torch.ops.higher_order.while_loop_stack_output,
    ):
        subgraph_args = (*args[2], *args[3])
        yield args[0], subgraph_args
        yield args[1], subgraph_args
    elif node.target is control_deps:
        assert not kwargs, (
            "Subgraph arguments can be renamed, so we cannot consistently "
            "handle kwargs at this point in the stack."
        )
        control_deps_subgraph = args[1]
        control_deps_args = tuple(args[2:])
        if control_deps_subgraph in valid_subgraphs:
            yield control_deps_subgraph, control_deps_args
        if isinstance(
            control_deps_subgraph, torch.fx.GraphModule
        ) and pytree.tree_any_only(
            torch.fx.GraphModule,
            lambda subgraph: subgraph in valid_subgraphs,
            control_deps_args,
        ):
            placeholder_to_arg = dict(
                zip(
                    control_deps_subgraph.graph.find_nodes(op="placeholder"),
                    control_deps_args,
                    strict=True,
                )
            )
            wrapped_nodes = [
                n for n in control_deps_subgraph.graph.nodes if n.op == "call_function"
            ]
            assert len(wrapped_nodes) == 1
            wrapped_node = wrapped_nodes[0]

            def replace_placeholder(item: Any) -> Any:
                if isinstance(item, torch.fx.Node):
                    return placeholder_to_arg.get(item, item)
                return item

            wrapped_args, wrapped_kwargs = pytree.tree_map(
                replace_placeholder,
                (wrapped_node.args, wrapped_node.kwargs),
            )
            yield from _extract_subgraphs_and_args(
                wrapped_node, valid_subgraphs, *wrapped_args, **wrapped_kwargs
            )
    else:
        warnings.warn(
            f"Please add support for subgraph args to function {node.target}!"
        )

        # By default, just return the detected list of subgraphs so that we can run
        # updates on all of them.
        flat_args_kwargs, _ = pytree.tree_flatten((args, kwargs))
        yield from (
            (s, None)
            for s in flat_args_kwargs
            if isinstance(s, torch.fx.GraphModule) and s in valid_subgraphs
        )


@dataclass(frozen=True)
class _FxNodeHash:
    node: torch.fx.Node
    target: torch.fx.node.Target
    arg_hashes: tuple[int, ...]
    kwarg_hashes: tuple[int, ...]


class FakeTensorUpdater:
    """
    The main idea here is that it's difficult to maintain accurate fake
    tensors (our primary form of metadata) for each node in our graph as we
    transform it.

    The most reliable way to obtain this information is by rerunning
    faketensor propagation. However, in general, faketensor propagation is
    fairly expensive. So, instead we'd like to only rerun faketensor
    propagation on nodes that have changed.

    In order to detect which nodes have changed, we first hash its node,
    target, and argument lists (which are immutable in FX).

    Then, whenever we call incremental_update, we check which FX nodes have a
    new hash, and recompute the faketensor metadata for that node. Then, we
    continue to recursively compute the faketensors for all users until the
    fake tensors stop changing.

    Since this runs in the context of Inductor, we assume that the input and
    output semantics for the outermost graph are not subject to change after class
    initialization, but we allow striding changes for subgraphs.  Any other changes will
    result in errors or undefined behavior.
    """

    def __init__(self, gm: torch.fx.GraphModule) -> None:
        self.processed_hashes = OrderedSet[_FxNodeHash]()
        self.gm = gm

        # Import here to avoid circular import issues.
        from torch._inductor.compile_fx import _get_subgraph_names

        self.subgraph_updaters: dict[torch.fx.GraphModule, FakeTensorUpdater] = {
            (subgraph := getattr(self.gm, subgraph_name)): FakeTensorUpdater(subgraph)
            for subgraph_name in _get_subgraph_names(self.gm)
        }

        for node in self.gm.graph.nodes:
            self.processed_hashes.add(self.hash_node(node))

    def hash_node(self, node: torch.fx.Node) -> _FxNodeHash:
        def get_hash_or_ids(n_iter: Iterable[Any]) -> tuple[int, ...]:
            """Replace unhashable items from the input with the ids of those items, and
            hash all other items.  This is kludgy, but allows us to attempt to account
            for unhashable classes like torch.fx.Node when hashing args and kwargs for
            other nodes."""

            def get_hash_or_id(o: object) -> int:
                try:
                    return hash(o)
                except Exception:
                    return id(o)

            return tuple(get_hash_or_id(n) for n in n_iter)

        return _FxNodeHash(
            node,
            node.target,
            get_hash_or_ids(node.args),
            get_hash_or_ids(chain.from_iterable(node.kwargs.items())),
        )

    def incremental_update(self) -> int:
        """Update FakeTensors on self.graph. We will try to do the minimum amount of work.

        Returns the number of nodes updated, including recursive updates on subgraphs."""
        existing_storages: defaultdict[int, int] = defaultdict(int)
        for node in self.gm.graph.nodes:
            if storage := get_node_storage(node):
                existing_storages[storage] += 1

        def should_process_node(node: torch.fx.Node) -> bool:
            return (
                callable(node.target)
                # node.target will called with FakeTensor arguments, which are not
                # supported by Inductor lowerings. TODO: Investigate how to remove
                # this. See https://github.com/pytorch/pytorch/issues/164920
                and not hasattr(node.target, "_inductor_lowering_function")
            )

        def node_invokes_subgraph(
            node: torch.fx.Node, *args: Any, **kwargs: Any
        ) -> bool:
            return (
                node.op == "call_function"
                and node.target
                not in (
                    # auto_functionalized doesn't call a subgraph, but the pytree call
                    # below can return subgraphs if the functionalized call itself
                    # invokes a subgraph.
                    torch.ops.higher_order.auto_functionalized,
                    torch.ops.higher_order.auto_functionalized_v2,
                )
                and pytree.tree_any_only(
                    torch.fx.GraphModule,
                    lambda s: s in self.subgraph_updaters,
                    (args, kwargs),
                )
            )

        # Update self.processed_hashes every time.  This allows us to account for
        # situations where a node gets modified, updated, then reverted to its original
        # state.  Without doing this, we wouldn't update downstream nodes, because the
        # reverted node would already be in our cache.
        current_graph_hashes = OrderedSet[_FxNodeHash]()

        nodes_updated: int = 0
        to_process = OrderedSet[int]()
        # Value records whether subgraph outputs have been updated.
        subgraph_updatings: dict[torch.fx.GraphModule, bool] = {}
        for node in self.gm.graph.nodes:
            current_graph_hashes.add(hash := self.hash_node(node))
            is_valid, args, kwargs = get_fake_args_kwargs(node, self.gm)
            if not is_valid:
                continue

            # NB: Be very careful about skipping nodes (via continues) here
            # and ask for a careful review when changing this code. The
            # consequence for incorrect FakeTensor metadata is difficult-to-debug
            # silent incorrectness.
            if (
                # Always run updates on nodes that invoke subgraphs
                not (invokes_subgraph := node_invokes_subgraph(node, *args, **kwargs))
                and hash in self.processed_hashes
                and id(node) not in to_process
            ):
                continue

            if not should_process_node(node):
                continue

            if invokes_subgraph:
                any_output_updated = False

                for subgraph, subgraph_args in _extract_subgraphs_and_args(
                    node, self.subgraph_updaters, *args, **kwargs
                ):
                    update_subgraph = subgraph not in subgraph_updatings
                    if update_subgraph:
                        subgraph_updatings[subgraph] = False

                    # If the arguments being passed into the subgraph differ from the
                    # existing args the first time we see the subgraph, update the
                    # placeholder nodes in the subgraph.  Any updates past that point
                    # would make subgraph updates inconsistent, so throw errors.
                    if subgraph_args is not None:
                        for p, a in zip(
                            subgraph.graph.find_nodes(op="placeholder"),
                            subgraph_args,
                            strict=True,
                        ):
                            # Do an update iff anything except the storage has changed,
                            # since every invocation of the subgraph will use different
                            # storages.  This implies potentially incorrect arg
                            # aliasing relationships.
                            if not _is_fake_tensor_same(
                                a,
                                p_fake := get_fake(p, subgraph),
                                existing_storages,
                                check_storage=False,
                            ):
                                assert update_subgraph, (
                                    "subgraph args must have consistent values!"
                                )
                                # Check that only the stride has changed.  Other changes
                                # cannot be handled without manual intervention.
                                assert _is_fake_tensor_same(
                                    a,
                                    p_fake,
                                    existing_storages,
                                    check_strides=False,
                                    check_storage=False,
                                ), (
                                    "A subgraph argument other than striding has been "
                                    "modified; FakeTensorUpdater cannot update this "
                                    "argument!"
                                )

                                p.meta["val"] = a
                                nodes_updated += 1

                    if update_subgraph:
                        _, orig_output_args, _ = get_fake_args_kwargs(
                            subgraph.graph.output_node(), subgraph
                        )
                        nodes_updated += self.subgraph_updaters[
                            subgraph
                        ].incremental_update()
                        _, new_output_args, _ = get_fake_args_kwargs(
                            subgraph.graph.output_node(), subgraph
                        )

                        if not _is_fake_tensor_same(
                            new_output_args,
                            orig_output_args,
                            existing_storages,
                        ):
                            subgraph_updatings[subgraph] = True

                    any_output_updated = (
                        any_output_updated or subgraph_updatings[subgraph]
                    )

                # If the outputs of the subgraphs have not changed (and have been
                # previously cached), we can skip updating.  If the outputs of the
                # subgraph have changed, we'll run FakeTensor propagation for every
                # invocation of the subgraph, so that each node gets unique FakeTensor
                # values which maintain appropriate aliasing relationships.
                if not any_output_updated and "val" in node.meta:
                    continue

            with V.fake_mode, enable_python_dispatcher():
                new_fake_tensor = node.target(*args, **kwargs)

            if "val" in node.meta and _is_fake_tensor_same(
                new_fake_tensor, node.meta["val"], existing_storages, node=node
            ):
                continue

            rebind_unbacked(V.fake_mode.shape_env, node, new_fake_tensor)

            node.meta["val"] = new_fake_tensor
            nodes_updated += 1

            if (shape_env := V.fake_mode.shape_env) and (
                symbol_to_path := compute_unbacked_bindings(shape_env, new_fake_tensor)
            ):
                # Refresh the bindings to the new symbols

                node.meta["unbacked_bindings"] = symbol_to_path

            if storage := get_node_storage(node):
                existing_storages[storage] += 1

            to_process.update(id(user) for user in node.users)

        self.processed_hashes = current_graph_hashes

        return nodes_updated


def get_storage(t: torch.Tensor) -> int:
    return t.untyped_storage()._cdata


def get_node_storage(node: torch.fx.Node) -> int | None:
    if "val" not in node.meta:
        return None
    if not isinstance(node.meta["val"], torch.Tensor):
        return None
    if not torch._C._has_storage(node.meta["val"]):
        return None
    return get_storage(node.meta["val"])


def get_fake(x: Any, gm: torch.fx.GraphModule | None) -> Any:
    """Return a fake tensor from the meta values of an input FX node.  If the input node
    is a get_attr node, we attempt to resolve it as a member of gm."""
    if isinstance(x, torch.fx.Node):
        if "val" in x.meta:
            return x.meta["val"]
        if "example_value" in x.meta:
            return x.meta["example_value"]
        if x.op == "get_attr" and isinstance(x.target, str) and hasattr(gm, x.target):
            return getattr(gm, x.target)
        # If a node has no available fake values and isn't a get_attr, it either returns
        # None or is incorrectly configured.  Since there are real nodes which return
        # None, we can't error here.
        return None
    # If there are no example values, return x
    return x


def get_fake_args_kwargs(
    x: torch.fx.Node, gm: torch.fx.GraphModule | None = None
) -> tuple[bool, tuple[Any, ...], dict[str, Any]]:
    """
    First value returns a boolean if any of the input nodes don't have a faketensor and
    weren't resolved from gm.
    """
    args, kwargs = tree_map(partial(get_fake, gm=gm), (x.args, x.kwargs))
    if any(
        isinstance(a, torch.fx.Node) for a in pytree.arg_tree_leaves(*args, **kwargs)
    ):
        return False, args, kwargs
    return True, args, kwargs


def is_node_realized(node: torch.fx.Node) -> bool:
    """Returns true if a node is always realized when lowered to inductor IR.

    NOTE: This may return some false negatives. e.g. it doesn't
    handle buffers realized heuristically during lowering, or
    buffers realized indirectly through view ops.
    """
    from torch._inductor.lowering import fallbacks, needs_realized_inputs

    def is_buffer(node: torch.fx.Node) -> bool:
        if node.op == "call_function" and node.target is operator.getitem:
            # For nodes with multiple outputs, we get the fx graph:
            #     foo = torch.ops.aten.foo(...)
            #     getitem = foo[0]
            #     getitem_1 = foo[1]
            # where we need to check if foo is a fallback kernel
            return is_buffer(node.args[0])  # type: ignore[arg-type]
        return node.op in ("placeholder", "output") or node.target in fallbacks

    if is_buffer(node):
        return True

    def realizes_inputs(node: torch.fx.Node) -> bool:
        return node.op == "output" or node.target in needs_realized_inputs

    if any(realizes_inputs(user) for user in node.users):
        return True

    # Otherwise, assume node isn't realized
    return False


def count_flops_fx(node: torch.fx.Node) -> int | None:
    if not countable_fx(node) or isinstance(node.target, str):
        return None
    with FakeTensorMode(allow_non_fake_inputs=True):
        success, args, kwargs = get_fake_args_kwargs(node)

        if success:
            # flex_attention HOPs have registered formulas, but invoking them
            # here can require tracing-only context, e.g. TransformGetItemToIndex.
            if node.target in (
                torch.ops.higher_order.flex_attention,
                torch.ops.higher_order.flex_attention_backward,
            ):
                flop_formula = flop_registry.get(node.target)
                if flop_formula is not None:
                    return flop_formula(*args, **kwargs, out_val=node.meta.get("val"))

            with torch.utils.flop_counter.FlopCounterMode(
                display=False
            ) as flop_counter_mode:
                node.target(*args, **kwargs)

            counted_flops = flop_counter_mode.get_total_flops()
            return counted_flops
    return None


def countable_fx(node: torch.fx.Node) -> bool:
    """
    Whether or not we can count the flops of an FX node.
    """
    assert isinstance(node, torch.fx.Node)
    if not hasattr(node, "target"):
        return False
    target = node.target
    if not hasattr(target, "overloadpacket"):
        return target in flop_registry
    packet = target.overloadpacket
    return packet in flop_registry