content: """Resource-aware graph partitioner for TensorRT compilation.
  
  This module refines an existing capability-based partitioning (accelerated vs
  non-accelerated subgraphs) by further splitting accelerated subgraphs to meet
  host CPU memory constraints during TensorRT engine building.
  
  High-level algorithm
  --------------------
  Given an original `torch.fx.GraphModule` and a capability-partitioned
  `GraphModule` (produced earlier in the pipeline), we:
  
  1) Reconstruct subgraphs on the original graph
     - Iterate over the capability-partitioned module to determine which original
       nodes belong to which subgraph (accelerated or not).
     - Preserve fusion groups discovered in each subgraph so that all nodes in a fusion
       group remain in the same subgraph and not be split across subgraphs.
     - Verify subgraphs respect topological order. This is to ensure the validity of the subgraphs.
     - Reconstruting subgraphs from partitioned module is easier than building nasted partitioned graph modules and flattening them later.
  
  2) Estimate memory cost of each possible subgraphs
     - Compute a per-subgraph "size" by traversing the graph to find weights
       (get_attr) reachable from its nodes and summing tensor bytes.
     - Use a set to record the visited nodes and avoid double counting shared parameters across subgraphs.
  
  
  4) Split large accelerated subgraphs
     - While a subgraph exceeds the per-engine budget, split it into two or more subgraphs.
     - Move nodes incrementally from the front of the original subgraph into a
       new left subgraph, repeatedly validating/correcting topological, partitioning, and
       dependency constraints.
     - Ensure we never split across a fusion group; when a split would break a
       fusion, we backtrack dependencies and move the entire fusion and related nodes into the left
       side.
     - Continue until the left subgraph fits the budget
     - Repeat the process for the right subgraph until all subgraphs fit the budget.
  
  5) Finalize
     - After splitting, assert all fusion groups reside in a single subgraph.
     - Tag nodes and produce a `GraphModule` where each subgraph becomes either a
       TRT engine (accelerated) or runs in Torch (non-accelerated).
  
  Notes
  -----
  - The budget is a heuristic bound. If the total model size exceeds 40x the
    per-engine budget, we fail early with a clear error suggesting remedies.
  """
  
  import logging
  from typing import Dict, List, Optional, Set, Tuple
  
  import psutil
  import torch
  from torch.fx.experimental.const_fold import _inline_module
  from torch.fx.passes.splitter_base import Subgraph, _SplitterBase
  from torch.fx.passes.tools_common import CALLABLE_NODE_OPS
  from torch_tensorrt.dynamo.partitioning._atomic_subgraphs import (
      get_node_in_fusion_pattern,
  )
  
  logger = logging.getLogger(__name__)
  
  MAX_NUM_OF_ENGINES = 50
  ENGINE_COMPILATION_MEMORY_USAGE_MULTIPLIER = 4
  
  
  class ResourcePartitioner(_SplitterBase):  # type: ignore
      """Refine capability-based subgraphs to meet host CPU memory constraints.
  
      This partitioner takes:
      - an original `torch.fx.GraphModule` (`module`)
      - a capability-partitioned `GraphModule` (`partitioned_module`) containing
        submodules that delineate accelerated vs non-accelerated regions
      - a CPU memory budget in bytes (`cpu_memory_budget`)
  
      It maps nodes from `module` into subgraphs according to `partitioned_module`
      and then splits oversized accelerated subgraphs so that each resulting TRT
      engine's estimated size fits within a conservative budget derived from
      available CPU memory or predefined CPU budget.
      """
  
      def __init__(
          self,
          module: torch.fx.GraphModule,
          cpu_memory_budget: Optional[int],
          submodule_name: str,
      ):
  
          assert isinstance(module, torch.fx.GraphModule)
  
          self.module = module
          used_rss: int = psutil.Process().memory_info().rss
          self.remaining_memory_budget = (
              cpu_memory_budget - used_rss
              if cpu_memory_budget is not None
              else psutil.virtual_memory().available
          )
          self.not_set_limit = cpu_memory_budget is None
          self.resource_split_count = 0
          self.submodule_name = submodule_name
          self.deps = self.find_deps()
  
          self._node_submodule_map: Dict[str, str] = {}
          self._return_tuple = False
          self.fusion_patterns: Dict[torch.fx.Node, Set[torch.fx.Node]] = {}
  
      def partition_graph(self) -> torch.fx.GraphModule:
          """Build the final partitioned `GraphModule` honoring memory constraints.
  
          Steps:
          - Build subgraph assignments from the capability-partitioned module
          - Split oversized accelerated subgraphs based on memory budget
          - Tag nodes and construct the final split graph
  
          Returns:
              torch.fx.GraphModule: A graph split into subgraphs based on capability partitioning and memory constraints.
          """
          # Delegate nodes based on operator coverage
          subgraphs = self.put_nodes_into_subgraphs()
          sizes = self.size_of_subgraphs(subgraphs)
          if (
              sum(sizes) * ENGINE_COMPILATION_MEMORY_USAGE_MULTIPLIER
              < self.remaining_memory_budget
          ):
              return self.module
  
          subgraphs = self.break_subgraphs(
              subgraphs, subgraph_size_budget=self.calculate_size_budget()
          )
  
          if len(subgraphs) == 1:
              return self.module
  
          # Set the number of TRT engines to be generated
          self.num_trt_accelerated_subgraphs = len([s for s in subgraphs if s.is_acc])
  
          # Tag the accelerated nodes and split the graph accordingly
          self.tag(subgraphs)
  
          gm = self.split(remove_tag=True)
  
          return gm
  
      def tag(self, subgraphs: list[Subgraph]) -> None:
          self.tags = []
          for subgraph in subgraphs:
              tag = f"{self.submodule_name}_resource_split_{self.resource_split_count}"
              self.resource_split_count += 1
              self.tags.append(tag)
              for node in subgraph.nodes:
                  node.tag = tag
                  self._node_submodule_map[node.name] = tag
  
      def put_nodes_into_subgraphs(self) -> list[Subgraph]:
          """
          Put the nodes into the subgraphs and erase the tag from previous partitioner if it exists.
          Returns:
              list[Subgraph]: Ordered subgraphs consisting of nodes in `module`.
          """
  
          nodes = []
          for node in self.module.graph.nodes:
              if hasattr(node, "tag"):
                  del node.tag
              if node.op in CALLABLE_NODE_OPS:
                  nodes.append(node)
          subgraphs = [Subgraph(is_acc=True, nodes=nodes)]
          self.fusion_patterns = get_node_in_fusion_pattern(self.module.graph)
  
          return subgraphs
  
      def check_topological_order(self, subgraphs: List[Subgraph]) -> bool:
          """Return True if subgraphs are in a valid topological order.
  
          Each node's dependencies must appear in earlier subgraphs or earlier
          positions within the same subgraph. Subgraphs should be topologically ordered to ensure the validity of the subgraphs.
          """
          visited_nodes: set[torch.fx.Node] = set()
          for subgraph in subgraphs:
              for node in subgraph.nodes:
                  if self.deps[node] > visited_nodes:
                      return False
                  visited_nodes.add(node)
          return True
  
      def calculate_size_budget(
          self,
          engine_compilation_memory_usage_multiplier: int = ENGINE_COMPILATION_MEMORY_USAGE_MULTIPLIER,
      ) -> int:
          """Compute the per-engine size budget in bytes.
  
          Uses explicit `cpu_memory_budget` minus used RSS
          divided by a safety multiplier.
  
          Args:
              engine_compilation_memory_usage_multiplier: Safety divisor applied to
                  available memory to approximate a per-engine budget. By default we assume TensorRT
                  compilation requires up to 4x the model's size.
  
          Returns:
              int: Budget in bytes for a single accelerated subgraph.
          """
  
          return (
              self.remaining_memory_budget // engine_compilation_memory_usage_multiplier
          )
  
      def break_subgraphs(
          self, subgraphs: List[Subgraph], subgraph_size_budget: int
      ) -> List[Subgraph]:
          """Split oversized accelerated subgraphs until they fit within budget.
  
          - Compute sizes for each subgraph (in bytes of parameters reachable from
            that subgraph).
          - If the sum of all sizes is catastrophically larger than budget
            (threshold 40x), raise a ValueError with guidance.
          - For any subgraph whose size exceeds `subgraph_size_budget`, iteratively
            split it using `break_subgraph_by_size` and append resulting segments.
          - Validate that fusion groups remain intact post splitting.
  
          Args:
              subgraphs: Ordered list of subgraphs from capability partitioning.
              subgraph_size_budget: Target maximum size per accelerated subgraph.
  
          Returns:
              List[Subgraph]: New list of subgraphs after resource-aware splitting.
          """
  
          new_subgraphs = []
          # We throw an error if the remaining memory is almost empty compared to the model size.
          # i.e. if the remaining memory is 4G (budget is 1G) the model size is greater than 40G, we stop the compilation.
          sizes = self.size_of_subgraphs(subgraphs)
          if sum(sizes) > subgraph_size_budget * MAX_NUM_OF_ENGINES:
              if self.not_set_limit:
                  raise ValueError(
                      "The system memory is too constrained to compile the model without severe perf degradation. Consider setting offload_module_to_cpu=False to save more CPU memory."
                  )
              else:
                  raise ValueError(
                      "CPU memory budget is too small to compile the model. "
                      + f"CPU memory budget: {self.remaining_memory_budget // (1024 * 1024)} MB, Model size: {sum(sizes) // (1024 * 1024)} MB. "
                      + "Consider setting cpu_memory_budget to a larger value."
                  )
          for subgraph, size in zip(subgraphs, sizes):
  
              while size > subgraph_size_budget:
                  broken_subgraphs, size_0, size_1 = self.break_subgraph_by_size(
                      subgraph, subgraph_size_budget
                  )
                  size = size_1
                  new_subgraphs.append(broken_subgraphs[0])
                  subgraph = broken_subgraphs[1]
  
              if len(subgraph.nodes) != 0:
                  new_subgraphs.append(subgraph)
  
          self._verify_all_fusion_nodes_in_same_subgraph(new_subgraphs)
  
          return new_subgraphs
  
      def _verify_all_fusion_nodes_in_same_subgraph(
          self, subgraphs: List[Subgraph]
      ) -> None:
          """Assert that every fusion group is contained in exactly one subgraph."""
          node_to_subgraph = {}
          for i, s in enumerate(subgraphs):
              for n in s.nodes:
                  node_to_subgraph[n] = i
  
          fusion_nodes_map_list = [
              len({node_to_subgraph[n] for n in ns}) == 1
              for ns in self.fusion_patterns.values()
          ]  # fusion nodes must be in the same subgraph
  
          assert all(
              fusion_nodes_map_list
          ), "All fusion nodes must be in the same subgraph"
          logger.info("All fusion nodes are in the same subgraph.")
  
      def break_subgraph_by_size(
          self, subgraph: Subgraph, size_to_break: int
      ) -> Tuple[List[Subgraph], int, int]:
          """Split a single oversized subgraph into two valid subgraphs.
  
          Moves nodes from the head of `subgraph` into a new left segment until
          the left segment's estimated size exceeds `size_to_break`. During the
          process we:
          - Repeatedly validate/correct topological placement
          - Detect and avoid splitting fusion groups by moving all fused nodes
            (and their producer chain) into the left segment
  
          Returns:
              (segments, size_left, size_right):
                  segments[0] is the new left subgraph, segments[1] is the residual
                  right subgraph. Sizes are estimated parameter bytes of each.
          """
          all_nodes = subgraph.nodes
          device_ordinal = subgraph.device_ordinal
          new_subgraphs = [
              Subgraph(
                  is_acc=True,
                  nodes=[],
                  device_ordinal=device_ordinal,
              ),
              Subgraph(
                  is_acc=True,
                  nodes=all_nodes,
                  device_ordinal=device_ordinal,
              ),
          ]
  
          # We break the subgraph until the left subgraph fits the budget.
          while True:
              # Set a step size proportional to the size of the subgraph to make the algorithm more efficient.
              # This reduce the time complexity from O(N**2) to O(N). The max number of steps is 50.
              # Note: we want the first step size to be 1.
              step_size = (
                  1 if not new_subgraphs[0].nodes else max(1, len(all_nodes) // 50)
              )
              new_subgraphs = self.step_and_validate(new_subgraphs, step_size)
              size_0, size_1 = self.size_of_subgraphs(new_subgraphs)
              if size_0 > size_to_break or size_0 > size_1:
                  break
  
          return new_subgraphs, size_0, size_1
  
      def step_and_validate(
          self, new_subgraphs: List[Subgraph], step_size: int = 1
      ) -> List[Subgraph]:
          """Advance the split by `step_size` nodes, then add more nodes to the left subgraph if rules are broken.
             There are two rules to check:
              1. The subgraphs should be ordered in a way that is safely to partition.
                 This is checked by validate_and_correct_subgraphs. Check that function for more details.
              2. The subgraphs should not break any fusion groups.
          - Move `step_size` nodes from the right to the left subgraph.
          - Run validation/correction to ensure a legal partitioning placement.
          - Get all leaf nodes in the left subgraph and check whether any of them are in a fusion group.
          - If the move splits a fusion group, migrate the entire fusion into the left subgraph.
  
          Returns:
              List[Subgraph]: Updated pair of subgraphs after stabilization.
          """
  
          for _ in range(step_size):
              new_subgraphs[0].nodes.append(new_subgraphs[1].nodes.pop(0))
  
          while True:
              new_subgraphs = self.validate_and_correct_subgraphs(new_subgraphs)
              nodes_in_first_subgraph = set(new_subgraphs[0].nodes)
              nodes_in_second_subgraph = set(new_subgraphs[1].nodes)
              leaf_node = self.get_leaf_node(nodes_in_first_subgraph)
              broken_fusion = self.step_if_break_fusion(
                  new_subgraphs,
                  leaf_node,
                  nodes_in_first_subgraph,
                  nodes_in_second_subgraph,
              )
              if not broken_fusion or len(new_subgraphs[1].nodes) == 0:
                  break
  
          return new_subgraphs
  
      def step_if_break_fusion(
          self,
          subgraphs: List[Subgraph],
          leaf_nodes: set[torch.fx.Node],
          nodes_in_first_subgraph: set[torch.fx.Node],
          nodes_in_second_subgraph: set[torch.fx.Node],
      ) -> bool:
          """Detect a fusion split and migrate fused nodes to the left subgraph.
  
          Given the current split boundary (captured by `leaf_nodes` of the left
          subgraph), check all recorded fusion groups. If any fused node remains
          on the right while its peer is on the left, pull the node and all of its
          producer chain into the left subgraph to keep fusions intact.
  
          Returns:
              bool: True if any fusion was migrated (i.e., a split would have
                    broken a fusion), otherwise False.
          """
  
          def add_nodes(node: torch.fx.Node) -> None:
              """
              This function adds a node and all its previous nodes to the first subgraph and removes it from the second subgraph in post order.
              """
              if (
                  node.op in CALLABLE_NODE_OPS
                  and node not in nodes_in_first_subgraph
                  and node in nodes_in_second_subgraph
              ):
                  # Exclude all nodes already in the first subgraph
                  nodes_in_first_subgraph.add(node)
                  nodes_in_second_subgraph.remove(node)
                  for input_node in node._input_nodes:
                      add_nodes(input_node)
                  subgraphs[0].nodes.append(node)
                  subgraphs[1].nodes.remove(node)
  
          fusion_broken = False
          for leaf in leaf_nodes:
              for node in self.fusion_patterns.get(leaf, []):
                  if (
                      node not in nodes_in_first_subgraph
                      and node in nodes_in_second_subgraph
                  ):
                      fusion_broken = True
                      add_nodes(node)
  
          return fusion_broken
  
      def get_leaf_node(
          self, nodes_in_first_subgraph: set[torch.fx.Node]
      ) -> set[torch.fx.Node]:
          """Return nodes in the left subgraph that feed any node on the right.
  
          A node is considered a leaf if at least one of its users is not in the
          left subgraph.
          """
          leaf_node = set()
  
          for node in nodes_in_first_subgraph:
              for user in node.users:
                  if user not in nodes_in_first_subgraph:
                      leaf_node.add(node)
                      break
          return leaf_node
  
      def size_of_subgraphs(self, subgraphs: List[Subgraph]) -> List[int]:
          """Estimate parameter footprint (bytes) for each subgraph.
  
          Traverses each subgraph's nodes and their producer chains to find
          parameters referenced via `get_attr`, summing tensor bytes. Shared
          parameters are counted only once globally.
  
          Returns:
              List[int]: Size per subgraph in bytes.
          """
          state_dict = self.module.state_dict(keep_vars=True)
          sizes = []
          weight_visited_nodes = set()
          for subgraph in subgraphs:
              nodes_in_subgraph = set(subgraph.nodes)
              stack = subgraph.nodes.copy()
              size = 0
              while stack:
                  node = stack.pop()
                  if node in weight_visited_nodes:
                      continue
                  weight_visited_nodes.add(node)
                  if node.op == "get_attr":
                      weight = state_dict.get(node.target, None)
                      if weight is None:
                          logger.warning(f"Weight {node.target} not found in state_dict")
                          continue
                      size += weight.numel() * weight.element_size()
                      continue
                  if node not in nodes_in_subgraph:
                      # Trace to other subgraphs
                      continue
                  for input_node in node._input_nodes:
                      if input_node not in weight_visited_nodes:
                          stack.append(input_node)
              sizes.append(size)
          return sizes
  
      def validate_and_correct_subgraphs(
          self, subgraphs: List[Subgraph]
      ) -> List[Subgraph]:
          """This is very important for the correctness of the partitioning. Torch gives undefined behavior if the subgraphs are not ordered correctly.
  
          The principle is: nodes that have all dependencies resolved in previous subgraphs should also be moved to the previous subgraph.
          For example, given a breakpoint node n resulting in two subgraphs S1 [..., n] and S2 [n+1, ...], all nodes in S2 that is not directly or indirectly depend on n should be moved to S1.
  
          We use a map to record the index of the subgraph that a node's users should belong to. If the node N is in subgraph S1 and is not the breakpoint node (subgraph.nodes[-1]),
          then the users that only depend on N should also be moved to S1. However, N is a breakpoint node, then the users that only depend on N should also be moved to S2.
  
          With the map, we can determine with subgraph a later node should be moved to according to all its inputs. We take max indices of all inputs nodes to determine the subgraph index.
  
          Returns:
              List[Subgraph]: Corrected subgraphs.
          """
          # a map from a node to the index of the subgraph it's user should belong to
          visited_nodes = {}
  
          for i, subgraph in enumerate(subgraphs):
              if i == 0:
                  for node in subgraph.nodes:
                      visited_nodes[node] = i
                  # breakpoint node's users should belong to the next subgraph
                  visited_nodes[subgraph.nodes[-1]] = i + 1
                  continue
  
              else:
                  to_remove_nodes = []
                  for j, node in enumerate(subgraph.nodes):
                      if j == len(subgraph.nodes) - 1:
                          # breakpoint node's users should belong to the next subgraph
                          visited_nodes[node] = i + 1
                          continue
                      subgraph_idx = 0
                      for dep in self.deps[node]:
                          if dep in visited_nodes:
                              # We take max indices of all inputs nodes to determine the subgraph index.
                              subgraph_idx = max(subgraph_idx, visited_nodes[dep])
  
                      if subgraph_idx != i:
                          # If the node should be moved to a different subgraph, we move it and remove it from the current subgraph.
                          subgraphs[subgraph_idx].nodes.append(node)
                          to_remove_nodes.append(node)
                      # Record the the subgraph that the users of this node should belong to
                      visited_nodes[node] = subgraph_idx
  
                  # Remove the nodes that are moved to other subgraphs
                  for node in to_remove_nodes:
                      subgraph.nodes.remove(node)
  
          return subgraphs
  
  
  def resource_partition(
      gm: torch.fx.GraphModule,
      cpu_memory_budget: int,
  ) -> torch.fx.GraphModule:
      """Resource-aware partitioning entry point.
  
      Takes an original FX graph (`gm`) and a capability-partitioned module
      (`partitioned_module`) and returns a final graph where accelerated segments
      are split further, if necessary, to satisfy CPU memory limits for TRT
      engine compilation.
  
      Args:
          gm: Original FX `GraphModule`.
          partitioned_module: Capability-partitioned `GraphModule` indicating
              accelerated vs non-accelerated regions.
          cpu_memory_budget: CPU memory budget in bytes for engine compilation.
              Use -1 to base the budget on currently available system memory.
  
      Returns:
          torch.fx.GraphModule: Final graph with resource-constrained subgraphs.
      """
  
      # Construct
      for name, _ in gm.named_children():
          submodule = getattr(gm, name)
          if (
              not isinstance(submodule, torch.fx.graph_module.GraphModule)
              or "_run_on_acc" not in name
          ):
              continue
          partitioner = ResourcePartitioner(
              submodule,
              submodule_name=name,
              cpu_memory_budget=cpu_memory_budget,
          )
  
          partitioned_graph = partitioner.partition_graph()
          setattr(gm, name, partitioned_graph)
  
      for name, module in list(gm.named_children()):
          split = False
          if "_run_on_acc" in name:
              for subname, submodule in module.named_children():
                  if "resource_split" in subname:
                      split = True
                      setattr(gm, subname, submodule)
              if split:
                  _inline_module(gm, name)
                  delattr(gm, name)
  
      gm.recompile()
      return gm
  