content: import logging
  from dataclasses import dataclass
  from typing import Collection, Dict, List, Optional, Tuple
  
  import torch
  import torch.fx.passes.operator_support as ops
  from torch.fx._compatibility import compatibility
  from torch.fx.node import Target
  from torch.fx.passes.splitter_base import (
      _SplitterBase,
      _SplitterSettingBase,
  )
  from torch.fx.passes.tools_common import (
      CALLABLE_NODE_OPS,
      FxNetAccFusionsFinder,
      NodeList,
      NodeSet,
      is_node_output_tensor,
  )
  from torch_tensorrt.dynamo._defaults import (
      MIN_BLOCK_SIZE,
      REQUIRE_FULL_COMPILATION,
  )
  from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
      DYNAMO_CONVERTERS as CONVERTERS,
  )
  from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
      ConverterRegistry,
  )
  
  logger = logging.getLogger(__name__)
  
  NON_COMPUTE_NODES = {"torch.ops.aten.view", "_operator.getitem"}
  NON_ACC_BACKEND_NAME = "None"
  
  
  @compatibility(is_backward_compatible=False)
  @dataclass
  class Subgraph:
      is_acc: bool
      backend: str
      nodes: NodeList
      device_ordinal: Optional[int] = None
  
  
  class BackendOpSupportTester(ops.OperatorSupportBase):  # type: ignore
      """Class to determine whether operators are supported by specific backends"""
  
      def __init__(
          self,
          backend_support_map: Dict[str, Collection[Target]],
          backend_priority: List[str],
          torch_executed_ops: Collection[Target] = set(),
      ) -> None:
          super().__init__()
  
          # Initialize sets of supported/unsupported operators
          self.supported_operators: Dict[str, int] = {}
          self.unsupported_operators: Dict[str, int] = {}
          self.torch_executed_ops = torch_executed_ops
          # Map of backend names to sets of supported operators
          self.backend_support_map = backend_support_map
          # Ordered list of backend names, from highest to lowest priority
          self.backend_priority = backend_priority
  
      def is_node_supported(
          self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node
      ) -> Tuple[bool, str]:
          node_name = ConverterRegistry.qualified_name_or_str(node.target)
  
          for i, backend_name in enumerate(self.backend_priority):
              supported_ops = self.backend_support_map.get(backend_name, set())
              supported_ops = {
                  ConverterRegistry.qualified_name_or_str(op) for op in supported_ops
              }
  
              if (
                  (node_name in supported_ops or node.op == "get_attr")
                  and node_name not in self.torch_executed_ops
                  and node.target not in self.torch_executed_ops
              ):
                  # If node is a proper, supported computational node, store the operator
                  if not node.is_impure() and node.op != "get_attr":
                      if node_name not in self.supported_operators:
                          self.supported_operators[f"{backend_name}_{node_name}"] = 1
                      else:
                          self.supported_operators[f"{backend_name}_{node_name}"] += 1
  
                  return True, backend_name
              else:
                  if i == len(self.backend_priority) - 1 and not node.is_impure():
                      if node_name not in self.unsupported_operators:
                          self.unsupported_operators[node_name] = 1
                      else:
                          self.unsupported_operators[node_name] += 1
  
          return False, NON_ACC_BACKEND_NAME
  
      def print_support_overview(self, num_acc_subgraphs: Optional[int] = None) -> None:
          if num_acc_subgraphs is not None:
              logger.debug(
                  f"\nNumber of Accelerated Subgraphs Generated: {num_acc_subgraphs}"
              )
  
          # Reformat support messages for debugger to print node overview as a single string
          supported_nodes_str = "\nSupported Nodes:\n"
          for node_name, count in self.supported_operators.items():
              supported_nodes_str += f"- {node_name} + Operator Count: {count}\n"
  
          logger.debug(supported_nodes_str)
  
          if self.unsupported_operators:
              unsupported_nodes_str = "\nUnsupported or Excluded Nodes:\n"
              for node_name, count in self.unsupported_operators.items():
                  unsupported_nodes_str += f"- {node_name} + Operator Count: {count}\n"
  
              logger.debug(unsupported_nodes_str)
          else:
              logger.debug("\nAll Nodes Supported\n")
  
  
  class HierarchicalAdjacencyPartitioner(_SplitterBase):  # type: ignore
      """Hierarchical Adjacency Partitioner to split an FX graph into subgraphs based on backend priority
  
      This partitioner extends the TRTPartitioner of adjacency_partitioner with backend priority awareness,
      allowing different parts of the model to be executed on different backends based on
      operator support and priority ordering.
  
      Args:
          module: FX GraphModule to partition
          operator_support: OperatorSupport class describing allowed operators
          backend_support_map: Dictionary mapping backend names to sets of supported operators
          backend_priority: Ordered list of backend names, from highest to lowest priority
          allowed_single_node_partition_ops: Nodes which can be included in single-node partitions
          min_block_size: Minimum number of computational operators per block
          require_full_compilation: Require that all computational operators be run in TRT
      Returns:
          torch.fx.GraphModule
      """
  
      def __init__(
          self,
          module: torch.fx.GraphModule,
          operator_support: ops.OperatorSupportBase,
          backend_support_map: Dict[str, Collection[Target]],
          backend_priority: List[str],
          allowed_single_node_partition_ops: Optional[Collection[str]] = None,
          min_block_size: int = MIN_BLOCK_SIZE,
          require_full_compilation: bool = REQUIRE_FULL_COMPILATION,
          return_tuple: bool = False,
          skip_fusion: bool = False,
      ):
          """
          Preprocesses graph before splitting with backend priority awareness
          """
          assert isinstance(module, torch.fx.GraphModule)
  
          self.module = module
          self.backend_support_map = backend_support_map
          self.backend_priority = backend_priority
  
          self.settings = _SplitterSettingBase(
              min_acc_module_size=min_block_size,
              allow_non_tensor=True,
              skip_fusion=skip_fusion,
          )
          self.operator_support = operator_support
  
          # Get all accelerated nodes based on operator support conditions
          self.acc_nodes = FxNetAccNodesFinder(
              self.module, self.operator_support, self.settings.allow_non_tensor
          )()
  
          if self.settings.skip_fusion:
              self.fusions = {}
          else:
              self.fusions = FxNetAccFusionsFinder(module, set(self.acc_nodes))()
  
          # Modify deps to add more deps for fused nodes
          self.deps = self.find_deps()
          self.update_deps_for_fusions()
  
          self.non_acc_submodule_name = "_run_on_gpu_"
          self._node_submodule_map: Dict[str, str] = {}
  
          self.num_accelerated_subgraphs: Optional[int] = None
          self.allowed_single_node_partition_ops = allowed_single_node_partition_ops
          self.require_full_compilation = require_full_compilation
          self._return_tuple = return_tuple
  
      def remove_small_acc_subgraphs(self, subgraphs: List[Subgraph]) -> List[Subgraph]:
          """
          This pass finds ACC submodules with less than specified size and merges
          them with adjacent GPU submodules.
          """
          result: List[Subgraph] = []
          for subgraph in subgraphs:
              if subgraph.is_acc:
                  if (
                      len(subgraph.nodes) >= self.settings.min_acc_module_size
                      or self.require_full_compilation
                      or (
                          self.allowed_single_node_partition_ops is not None
                          and any(
                              ConverterRegistry.qualified_name_or_str(node.target)
                              in self.allowed_single_node_partition_ops
                              for node in subgraph.nodes
                          )
                      )
                  ):
                      result.append(subgraph)
                  else:
                      logger.debug(
                          "Eliminating acc subgraph because it's smaller than the threshold: "
                          f"{len(subgraph.nodes)} < {self.settings.min_acc_module_size}"
                      )
                      # if the last subgraph result[-1] is non-acc or has the same backend, merge the current subgraph into it
                      if result and (
                          not result[-1].is_acc or result[-1].backend == subgraph.backend
                      ):
                          result[-1].nodes.extend(subgraph.nodes)
                      else:
                          # if the last subgraph result[-1] has different backends, then append the current subgraph as non-acc
                          subgraph.is_acc = False
                          subgraph.backend = "None"
                          result.append(subgraph)
              else:
                  if result and not result[-1].is_acc:
                      result[-1].nodes.extend(subgraph.nodes)
                  else:
                      if result:
                          if result[-1].backend == subgraph.backend:
                              result[-1].nodes.extend(subgraph.nodes)
                          else:
                              result.append(subgraph)
                      else:
                          result.append(subgraph)
          return result
  
      def partition_graph(self) -> torch.fx.GraphModule:
          """Partitions the GraphModule into subgraphs based on backend priority
  
          Returns a GraphModule with submodules for each segment
          """
          # Delegate nodes based on operator coverage
          subgraphs = self.put_nodes_into_subgraphs()
  
          # A graph is fully supported if there is a single partition and all operators are supported/convertible
          full_support = len([s for s in subgraphs if s.is_acc]) == 1 and not getattr(
              self.operator_support, "unsupported_operators", True
          )
  
          if not full_support and self.require_full_compilation:
              raise AssertionError(
                  "require_full_compilation=True was specified, but model is not fully supported or multiple partitions are found"
              )
  
          if (
              full_support
              and self.require_full_compilation
              and self.settings.min_acc_module_size != MIN_BLOCK_SIZE
          ):
              logger.warning(
                  "Detected both require_full_compilation and min_block_size compilation "
                  "arguments were specified. Disregarding min_block_size argument for "
                  "fully supported model."
              )
  
          # Remove segments smaller than the block size (with exceptions)
          subgraphs = self.remove_small_acc_subgraphs(subgraphs)
  
          # Set the number of accelerated subgraphs to be generated
          self.num_accelerated_subgraphs = len([s for s in subgraphs if s.is_acc])
  
          # Tag the accelerated nodes and split the graph accordingly
          self.tag(subgraphs)
          return self.split()
  
      def starter_nodes(self) -> Tuple[NodeSet, NodeSet]:
          """Generates starter nodes for partitioning + segmentation"""
          # Starter accelerated nodes are all callable accelerated ops
          starter_acc_nodes = {
              node for node in self.acc_nodes if node.op in CALLABLE_NODE_OPS
          }
  
          # Started non-accelerated nodes are the rest of the callable nodes
          starter_non_acc_nodes = {
              node
              for node in self.module.graph.nodes
              if (node not in starter_acc_nodes and node.op in CALLABLE_NODE_OPS)
          }
          return starter_non_acc_nodes, starter_acc_nodes
  
      def put_nodes_into_subgraphs(self) -> list[Subgraph]:
          # We start graph traversal from leaf nodes
          current_cpu_nodes, current_acc_nodes = self.starter_nodes()
          visited_nodes: NodeSet = set()
  
          # Determine which subgraph to start from based on which subgraph has
          # 0-dep node
          acc_subgraph: bool = not any(len(self.deps[n]) == 0 for n in current_cpu_nodes)
  
          current_subgraph_nodes: NodeList = []
  
          # Result accumulator
          subgraphs: list[Subgraph] = []
          while current_cpu_nodes or current_acc_nodes:
              # Find the first node that should belong to the current subgraph and has all dependencies resolved
              current_nodes = current_acc_nodes if acc_subgraph else current_cpu_nodes
              node = next(
                  (n for n in current_nodes if self.deps[n] <= visited_nodes),
                  None,
              )
  
              # If no node was found, then it's time to flip the mode and start a new subgraph
              if node is None:
                  if not current_subgraph_nodes:
                      raise FxNetSplitterInternalError("Subgraph can't be empty")
  
                  subgraphs.append(
                      Subgraph(
                          is_acc=acc_subgraph,
                          nodes=current_subgraph_nodes,
                          backend=(
                              current_subgraph_nodes[-1].backend
                              if acc_subgraph
                              else "None"
                          ),
                      )
                  )
                  acc_subgraph = not acc_subgraph
                  current_subgraph_nodes = []
                  continue
  
              # If the backend changed, then it's time to start a new subgraph
              if (
                  current_subgraph_nodes
                  and current_subgraph_nodes[-1].backend != node.backend
              ):
                  if not current_subgraph_nodes:
                      raise FxNetSplitterInternalError("Subgraph can't be empty")
  
                  subgraphs.append(
                      Subgraph(
                          is_acc=acc_subgraph,
                          nodes=current_subgraph_nodes,
                          backend=current_subgraph_nodes[-1].backend,
                      )
                  )
                  current_subgraph_nodes = []
                  continue
  
              current_nodes.remove(node)
              visited_nodes.add(node)
              current_subgraph_nodes.append(node)
  
              # Add fusion buddies
              if node in self.fusions:
                  if node in self.acc_nodes:
                      current_acc_nodes.update(self.fusions[node] - visited_nodes)
                  else:
                      current_cpu_nodes.update(self.fusions[node] - visited_nodes)
  
              # Put depending nodes into the queue
              for user in node.users:
                  if user.op not in CALLABLE_NODE_OPS:
                      continue
  
                  # Add downstream nodes
                  if user in self.acc_nodes:
                      current_acc_nodes.add(user)
                  else:
                      current_cpu_nodes.add(user)
  
          # Check if the last subgraph was not created
          if current_subgraph_nodes:
              subgraphs.append(
                  Subgraph(
                      is_acc=acc_subgraph,
                      nodes=current_subgraph_nodes,
                      backend=(
                          current_subgraph_nodes[-1].backend if acc_subgraph else "None"
                      ),
                  )
              )
  
          if not subgraphs:
              raise FxNetSplitterInternalError("Couldn't create subgraphs")
  
          return subgraphs
  
      def tag(self, subgraphs: list[Subgraph]) -> None:
          self.tags: list[str] = []
          for subgraph in subgraphs:
              tag = (
                  f"_run_on_acc_{subgraph.backend}_{len(self.tags)}"
                  if subgraph.is_acc
                  else f"{self.non_acc_submodule_name}{len(self.tags)}"
              )
              self.tags.append(tag)
              for node in subgraph.nodes:
                  if hasattr(node, "tag"):
                      raise FxNetSplitterInternalError(f"Node {node} was already tagged")
  
                  node.tag = tag
                  self._node_submodule_map[node.name] = tag
  
  
  @compatibility(is_backward_compatible=False)
  class FxNetAccNodesFinder:
      """
      Finds a set of nodes that can be supported on ACC, excluding nodes that have non-tensor
      input/output to cpu nodes to prevent non-tensor data flow between backends and cpu.
  
      I.e. if we have a chain:
  
      ACC_NODE_1 -> ACC_NODE_2 -> ACC_NODE_3 -> CPU_NODE_1
  
      where every ACC node produces non-tensor output, then they all should be treated as CPU nodes.
  
      This behavior can be turned off by passing allow_non_tensor=True.
      """
  
      def __init__(
          self,
          module: torch.fx.GraphModule,
          operator_support: ops.OperatorSupportBase,
          allow_non_tensor: bool,
      ):
          self.module = module
          self.operator_support = operator_support
          self.allow_non_tensor = allow_non_tensor
          self.acc_nodes: NodeSet = set()
  
      def reduce_acc_nodes_non_tensor_input_helper(self, cpu_worklist: NodeList) -> None:
          """
          Transitively excludes nodes from ACC supported set.
          For every node in the worklist:
          - removes its downstream ACC nodes from ACC supported set,
          - if any downstream ACC node produces non-tensor output,
            then it gets added into the worklist.
          """
          while cpu_worklist:
              node = cpu_worklist.pop(0)
  
              for user in node.users:
                  if user in self.acc_nodes:
                      self.acc_nodes.remove(user)
                      if not is_node_output_tensor(user):
                          cpu_worklist.append(user)
  
      def reduce_acc_nodes_non_tensor_input(self) -> None:
          """
          Excludes nodes from ACC supported set that have direct
          upstream CPU nodes that produce non-tensor outputs.
          """
          non_tensor_cpu_nodes: NodeList = []
  
          for node in self.module.graph.nodes:
              if node.op not in CALLABLE_NODE_OPS:
                  continue
              if node in self.acc_nodes:
                  continue
              if is_node_output_tensor(node):
                  continue
              non_tensor_cpu_nodes.append(node)
  
          self.reduce_acc_nodes_non_tensor_input_helper(non_tensor_cpu_nodes)
  
      def reduce_acc_nodes_non_tensor_output(self) -> None:
          """
          Excludes nodes from ACC supported set that produce non-tensor
          outputs and have downstream CPU nodes.
          """
          while True:
              new_cpu_nodes: NodeList = []
  
              for acc_node in self.acc_nodes:
                  if is_node_output_tensor(acc_node):
                      continue
                  for user in acc_node.users:
                      if user not in self.acc_nodes:
                          new_cpu_nodes.append(acc_node)
                          break
  
              if not new_cpu_nodes:
                  break
  
              for new_cpu_node in new_cpu_nodes:
                  self.acc_nodes.remove(new_cpu_node)
  
              self.reduce_acc_nodes_non_tensor_input_helper(new_cpu_nodes)
  
      def __call__(self) -> NodeSet:
          submodules = dict(self.module.named_modules())
          backend = NON_ACC_BACKEND_NAME
          for n in self.module.graph.nodes:
              # Group non-compute nodes with previous compute nodes
              if ConverterRegistry.qualified_name_or_str(n.target) in NON_COMPUTE_NODES:
                  n.backend = backend
                  if backend != NON_ACC_BACKEND_NAME:
                      self.acc_nodes.add(n)
                  continue
  
              if n.op in CALLABLE_NODE_OPS:
                  is_supported, backend = self.operator_support.is_node_supported(
                      submodules, n
                  )
                  if is_supported:
                      n.backend = backend
                      self.acc_nodes.add(n)
                  else:
                      n.backend = NON_ACC_BACKEND_NAME
  
          if not self.allow_non_tensor:
              self.reduce_acc_nodes_non_tensor_input()
              self.reduce_acc_nodes_non_tensor_output()
  
          return self.acc_nodes
  
  
  @compatibility(is_backward_compatible=False)
  class FxNetSplitterInternalError(Exception):
      pass
  
  
  def hierarchical_adjacency_partition(
      gm: torch.fx.GraphModule,
      min_block_size: int = MIN_BLOCK_SIZE,
      torch_executed_ops: Collection[Target] = set(),
      backend_support_map: Optional[Dict[str, Collection[Target]]] = None,
      backend_priority: Optional[List[str]] = None,
      require_full_compilation: bool = REQUIRE_FULL_COMPILATION,
      skip_fusion: bool = False,
  ) -> Tuple[torch.fx.GraphModule, BackendOpSupportTester]:
      """Partition an FX GraphModule with aten ops into submodules using hierarchical partitioning
      based on backend priority and operator support
  
      Args:
          gm: FX GraphModule to partition
          min_block_size: Minimum number of operators per TRT-Engine Block
          backend_support_map: Dictionary mapping backend names to sets of supported operators
          backend_priority: Ordered list of backend names, from highest to lowest priority
          require_full_compilation: Require that all computational operators be run in TRT
          skip_fusion: Skip fusions found by FxNetAccFusionsFinder
      Returns:
          torch.fx.GraphModule, BackendOpSupportTester
      """
      # Ensure graph is clean prior to partitioning
      gm.graph.eliminate_dead_code()
      gm.graph.lint()
      gm.recompile()
  
      # Default backend support map if none provided
      if backend_support_map is None:
          backend_support_map = {
              "tensorrt": CONVERTERS.keys(),
              "inductor": set(),
          }
  
      # Default backend priority if none provided
      if backend_priority is None:
          backend_priority = ["tensorrt", "inductor"]
  
      # Construct BackendOpSupportTester
      supported_ops = BackendOpSupportTester(
          backend_support_map=backend_support_map,
          backend_priority=backend_priority,
          torch_executed_ops=torch_executed_ops,
      )
      partitioner = HierarchicalAdjacencyPartitioner(
          gm,
          supported_ops,
          backend_support_map=backend_support_map,
          backend_priority=backend_priority,
          min_block_size=min_block_size,
          require_full_compilation=require_full_compilation,
          skip_fusion=skip_fusion,
      )
  
      partitioned_graph = partitioner.partition_graph()
  
      supported_ops.print_support_overview(partitioner.num_accelerated_subgraphs)
  
      return partitioned_graph, supported_ops
  