content: from collections import defaultdict
  from functools import lru_cache
  from typing import Any, Callable, Dict, List, Set, Tuple
  
  import torch
  from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
  from torch.ops import aten
  
  ATOMIC_SUBGRAPHS = []
  
  
  def register_atomic_subgraph(
      init_args: Tuple[Any, ...] = tuple(),
      is_core_aten: bool = False,
  ) -> Callable[[torch.nn.Module], torch.nn.Module]:
  
      def decorator(subgraph: torch.nn.Module) -> torch.nn.Module:
          ATOMIC_SUBGRAPHS.append((subgraph, init_args, is_core_aten))
          return subgraph
  
      return decorator
  
  
  @register_atomic_subgraph(init_args=(aten.silu.default,), is_core_aten=True)
  @register_atomic_subgraph(init_args=(aten.gelu.default,), is_core_aten=True)
  @register_atomic_subgraph(init_args=(aten.relu.default,), is_core_aten=True)
  @register_atomic_subgraph(init_args=(aten.sigmoid.default,), is_core_aten=True)
  class ConvBNActivation(torch.nn.Module):  # type: ignore[misc]
      def __init__(self, activation: torch._ops.OpOverload) -> None:
          super().__init__()
          self.activation = activation
  
      def forward(
          self,
          x: torch.Tensor,
          weight: torch.Tensor,
          bias: torch.Tensor,
          stride: List[int],
          padding: List[int],
          dilation: List[int],
          transposed: bool,
          output_padding: List[int],
          groups: int,
          bn_weight: torch.Tensor,
          bn_bias: torch.Tensor,
          running_mean: torch.Tensor,
          running_var: torch.Tensor,
          momentum: float,
          eps: float,
      ) -> torch.Tensor:
          x = aten.convolution.default(
              x,
              weight,
              bias,
              stride,
              padding,
              dilation,
              transposed,
              output_padding,
              groups,
          )
          x = aten._native_batch_norm_legit_no_training.default(
              x, bn_weight, bn_bias, running_mean, running_var, momentum, eps
          )[0]
          x = self.activation(x)
          return x
  
  
  @register_atomic_subgraph(init_args=(aten.silu.default,), is_core_aten=True)
  @register_atomic_subgraph(init_args=(aten.gelu.default,), is_core_aten=True)
  @register_atomic_subgraph(init_args=(aten.relu.default,), is_core_aten=True)
  @register_atomic_subgraph(init_args=(aten.sigmoid.default,), is_core_aten=True)
  class ConvActivation(torch.nn.Module):  # type: ignore[misc]
      def __init__(self, activation: torch._ops.OpOverload) -> None:
          super().__init__()
          self.activation = activation
  
      def forward(
          self,
          x: torch.Tensor,
          weight: torch.Tensor,
          bias: torch.Tensor,
          stride: List[int],
          padding: List[int],
          dilation: List[int],
          transposed: bool,
          output_padding: List[int],
          groups: int,
      ) -> torch.Tensor:
          x = aten.convolution.default(
              x,
              weight,
              bias,
              stride,
              padding,
              dilation,
              transposed,
              output_padding,
              groups,
          )
          x = self.activation(x)
          return x
  
  
  @register_atomic_subgraph(init_args=(), is_core_aten=True)
  class MulAdd(torch.nn.Module):  # type: ignore[misc]
      def __init__(self) -> None:
          super().__init__()
  
      def forward(
          self, x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
      ) -> torch.Tensor:
          x = aten.mul.Tensor(x, weight)
          x = aten.add.Tensor(x, bias)
          return x
  
  
  @register_atomic_subgraph(init_args=(), is_core_aten=True)
  class MulMul(torch.nn.Module):  # type: ignore[misc]
      def __init__(self) -> None:
          super().__init__()
  
      def forward(
          self, x: torch.Tensor, y: torch.Tensor, z: torch.Tensor
      ) -> torch.Tensor:
          x = aten.mul.Tensor(x, y)
          x = aten.mul.Tensor(x, z)
          return x
  
  
  def get_node_in_fusion_pattern(
      graph: torch.fx.Graph,
  ) -> Dict[torch.fx.Node, Set[torch.fx.Node]]:
      """
      This function gets the nodes map of the fusion pattern from the graph.
      Key: node that appears in the fusion pattern
      Value: the list of nodes that should be fused together
      """
      fusion_nodes = defaultdict(set)
      for compiled_pattern_graph in get_compiled_atomic_subgraphs():
          subgraph_matcher = SubgraphMatcher(compiled_pattern_graph.graph)
          match_result = subgraph_matcher.match(graph)
          for match in match_result:
              fusion_group = {
                  node
                  for node in match.nodes_map.values()
                  if node
                  and type(node) == torch.fx.Node
                  and node.op == "call_function"
                  and node not in match.placeholder_nodes
              }
              for node in fusion_group:
                  fusion_nodes[node].update(fusion_group)
  
      return fusion_nodes
  
  
  def get_compiled_atomic_subgraphs() -> List[torch.fx.GraphModule]:
      """
      This function gets the compiled atomic subgraphs from the graph.
      LRU cache the result to avoid recompiling the same pattern multiple times.
      """
      compiled_atomic_subgraphs = []
      for pattern, init_args, is_core_aten in ATOMIC_SUBGRAPHS:
          pattern_graph = trace_atomic_graph(pattern, init_args, is_core_aten)
          if not is_core_aten:
              # TODO: Add decomposition and lowering if is_core_aten is False
              raise NotImplementedError(
                  "Atomic subgraphs are not supported for non-aten subgraphs yet."
              )
          compiled_atomic_subgraphs.append(pattern_graph)
      return compiled_atomic_subgraphs
  
  
  @lru_cache(maxsize=None)
  def trace_atomic_graph(
      graph: torch.nn.Module, init_args: Any, is_core_aten: bool = True
  ) -> torch.fx.GraphModule:
      if is_core_aten:
          return torch.fx.symbolic_trace(graph(*init_args))
      else:
          raise NotImplementedError(
              "Resource partitioner currently does not support unlowered atomic subgraphs"
          )
  