code: def reorder_compute_for_overlap(
      snodes: List["scheduler.BaseSchedulerNode"],
  ) -> List["scheduler.BaseSchedulerNode"]:
      """
      Decides a global ordering of all compute and communication nodes,
      assuming that we already have a global ordering of communication nodes.
  
      Overall scheduling procedure is:
          Step 1: Given that we've currently scheduled comm N, we now schedule all compute nodes
              that are required for comm N + 1 but do not depend on comm N, to run at the same time with comm N.
          Step 2: If all those compute nodes are sufficient to overlap comm N, we're done.
              Otherwise, we now need to look elsewhere to find compute that overlaps with comm N.
              We prioritize compute nodes that are needed sooner.
          Step 3: We schedule the compute nodes dependent on comm N and required for comm N + 1.
          Step 4: We schedule comm N + 1.
          Repeat this for subsequent comm nodes.
      """
      final_order = []
  
      comm_nodes = []
      for snode in snodes:
          if isinstance(snode.node, ir.CollectiveKernel):
              comm_nodes.append(snode)
      if len(comm_nodes) == 0:
          # if there is no comm nodes, return the current order
          return snodes
  
      comm_ancestors = {
          node: get_ancestors(node) for node in comm_nodes
      }
      comm_descendants = {
          node: get_descendants(node) for node in comm_nodes
      }
  
      indeg = {k: 0 for k in snodes}
      for snode in snodes:
          for user in snode.node_users:
              if user in indeg:
                  indeg[user] += 1
      ready_to_schedule_nodes = {
          node for node in snodes if indeg[node] == 0
      }
  
      unscheduled_nodes = set()
      unscheduled_nodes = set(snodes)
  
      def schedule_node(snode):
          """
          Schedule a single node.
          """
          assert snode in unscheduled_nodes
          assert snode in ready_to_schedule_nodes
          ready_to_schedule_nodes.remove(snode)
          unscheduled_nodes.remove(snode)
          final_order.append(snode)
          for user in tuple_sorted(snode.node_users):
              if user in indeg:
                  indeg[user] -= 1
                  if indeg[user] == 0:
                      ready_to_schedule_nodes.add(user)
  
      def schedule_nodes(snodes):
          """
          Schedules all nodes in `snodes` in an arbitrary topologically valid order.
          """
          all_nodes = set(snodes)
          assert all(node in unscheduled_nodes for node in all_nodes)
          while len(all_nodes) > 0:
              # NOTE: since model graph is always a DAG and does not have circular dependency inside,
              # there should be at least one node that is a "free node" (i.e. indeg == 0),
              # hence infinite loop is not possible. But we check here just to be safe.
              progress = False
              for node in tuple_sorted(all_nodes):
                  if node in ready_to_schedule_nodes:
                      schedule_node(node)
                      all_nodes.remove(node)
                      progress = True
              if not progress:
                  raise Exception(
                      "Unable to find a free node (indeg == 0). This is an impossible state to reach. "
                      "Please report a bug to PyTorch."
                  )
  
      # First, schedule all compute nodes that are required by first comm node,
      # as well as the first comm node itself.
      assert len(comm_nodes) > 0
      schedule_nodes(
          list(comm_ancestors[comm_nodes[0]]) + [comm_nodes[0]],
      )
  
      rolled_over_compute_cost = 0
      for idx in range(1, len(comm_ancestors)):
          # Step 1: Given that we've currently scheduled comm `idx-1`, we now schedule
          # all compute nodes that are required for comm `idx` but do not depend on comm `idx-1`,
          # to run at the same time with comm `idx-1`.
          needed_by_next_comm_and_ready_compute_nodes = (
              unscheduled_nodes
              & (
                  comm_ancestors[comm_nodes[idx]]
                  - comm_descendants[comm_nodes[idx - 1]]
              )
          )
          assert_no_comm_nodes(
              needed_by_next_comm_and_ready_compute_nodes
          )
  
          total_compute_runtime_cost = rolled_over_compute_cost + sum(
              [
                  estimate_op_runtime(node)
                  for node in needed_by_next_comm_and_ready_compute_nodes
              ]
          )
          prev_comm_runtime_cost = estimate_op_runtime(
              comm_nodes[idx - 1]
          )
          schedule_nodes(
              tuple_sorted(needed_by_next_comm_and_ready_compute_nodes)
          )
  
          # Step 2: If all those compute nodes are sufficient to overlap comm `idx-1`, we're done.
          # Otherwise, we now need to look elsewhere to find compute that overlaps with comm `idx`.
          # We prioritize compute nodes that are needed sooner.
          step1_runtime_cost = total_compute_runtime_cost
          if step1_runtime_cost >= prev_comm_runtime_cost:
              pass
          else:
              # Find all ready to schedule compute nodes that do not depend on comm `idx-1`.
              ready_to_schedule_compute_nodes = tuple_sorted(
                  ready_to_schedule_nodes
                  - comm_descendants[comm_nodes[idx - 1]]
              )
              assert_no_comm_nodes(ready_to_schedule_compute_nodes)
  
              def earliest_comm_descendant(node):
                  for idx in range(len(comm_nodes)):
                      if node in comm_ancestors[comm_nodes[idx]]:
                          return idx
                  return len(comm_nodes)
  
              # Prioritize compute nodes that are needed sooner.
              ready_to_schedule_compute_nodes = sorted(
                  ready_to_schedule_compute_nodes,
                  key=earliest_comm_descendant,
              )
  
              for snode in ready_to_schedule_compute_nodes:
                  if (
                      total_compute_runtime_cost
                      >= prev_comm_runtime_cost
                  ):
                      # If accumulated compute runtime cost is greater than comm `idx-1` runtime cost,
                      # it means we have maximized overlap for comm `idx-1`, and hence we stop looking
                      # for more compute to schedule.
                      break
                  compute_runtime_cost = estimate_op_runtime(snode)
                  # If we're not able to leverage more than half of this
                  # node's compute to overlap, we skip it.
                  # TODO: Smarter heuristics here
                  if (
                      prev_comm_runtime_cost
                      - total_compute_runtime_cost
                  ) <= compute_runtime_cost / 2:
                      continue
                  schedule_node(snode)
                  total_compute_runtime_cost += compute_runtime_cost
          rollable_compute_cost = (
              total_compute_runtime_cost - step1_runtime_cost
          )
  
          # Step 3: We schedule the compute nodes dependent on comm `idx-1` and required for comm `idx`.
          needed_by_next_comm_nodes = (
              unscheduled_nodes & comm_ancestors[comm_nodes[idx]]
          )
          schedule_nodes(list(needed_by_next_comm_nodes))
  
          # Step 4: We schedule comm `idx`.
          schedule_nodes([comm_nodes[idx]])
  
          is_prev_comm_blocking_next_comm = (
              len(needed_by_next_comm_nodes) > 0
          )
          # The idea here is that if there are no compute nodes from Step 3
          # (i.e. if prev comm is not blocking next comm), we can roll over the compute nodes
          # in Step 2 to overlap with the next comm, since they're not required to finish
          # before the next comm starts.
          if is_prev_comm_blocking_next_comm:
              rolled_over_compute_cost = 0
          else:
              rolled_over_compute_cost = rollable_compute_cost  # type: ignore[assignment]
  
      schedule_nodes(unscheduled_nodes)
      return final_order
  
file_path: pytorch/torch/_inductor/comms.py
imports: from typing import List
  import torch
  from .None import config, ir, scheduler
  from .dependencies import WeakDep
  from .utils import tuple_sorted
label_desc: Use image node_id 8 for calling a global function with example usage: reorder_compute_for_overlap(snodes) and returns: final_order, snodes, len, idx