import operator from collections.abc import Callable from typing import Any, TYPE_CHECKING, TypeAlias import torch.fx import torch.fx.traceback import torch.utils._pytree as pytree from torch._dynamo.graph_utils import _get_flat_args from torch._dynamo.variables.streams import get_current_stream, new_event from torch.fx.node import map_arg from torch.utils._ordered_set import OrderedSet from torch.utils._runtime_estimation import ( _FLOAT_TYPES, _IGNORE_OPS, get_compute_time, get_transfer_time, ) if TYPE_CHECKING: from .schemas import ViewAndMutationMeta from .indexed_dict import IndexedDict aten = torch.ops.aten Node: TypeAlias = torch.fx.Node Graph: TypeAlias = torch.fx.Graph Partition: TypeAlias = bool FORWARD: Partition = False BACKWARD: Partition = True _EMPTY_LOCATIONS: frozenset[tuple[Partition, int]] = frozenset() _SYNC_OPS = ( torch.ops.streams.record_event.default, torch.ops.streams.wait_event.default, torch.ops.streams.wait_stream.default, torch.ops.streams.synchronize_event.default, torch.ops.streams.synchronize_device.default, torch.ops.streams.synchronize_stream.default, ) def get_roofline_estimate(node: Node) -> float: if node.op != "call_function": raise AssertionError(f"non-func node in roofline estimate: {node.op}") def map_value(x: Any) -> Any: return x.meta.get("value", x) if isinstance(x, Node) else x func = node.target if func in _IGNORE_OPS: return 0.0 mapped_args = torch.fx.map_arg(node.args, map_value) mapped_kwargs = torch.fx.map_arg(node.kwargs, map_value) flat_args_kwargs = [map_value(x) for x in _get_flat_args(node, {})] flat_outs, _ = pytree.tree_flatten(node.meta.get("value", node)) out = node.meta.get("value", node) out_dtypes = { t.dtype for t in flat_outs if isinstance(t, torch.Tensor) and t.dtype in _FLOAT_TYPES } return ( max( get_transfer_time(flat_args_kwargs, flat_outs), get_compute_time(func, mapped_args, mapped_kwargs, out, out_dtypes), ) / 1e6 ) def is_gradient_acc(node: Node) -> bool: return node.meta.get("is_gradient_acc", False) def is_bwd_node(node: Node) -> Partition: tag = node.meta.get("partitioner_tag") return tag == "is_backward" or tag == "must_be_in_backward" def get_device(node: Node) -> torch.device: return node.meta["val"].device def get_stream(node: Node) -> int | None: maybe_annotation = node.meta.get("custom", None) if maybe_annotation is not None: return node.meta["custom"].get("stream", None) else: return None def get_stream_or_current_stream(node: Node) -> int: ind = get_stream(node) if ind is None: ind = get_current_stream(get_device(node)) return ind def set_stream(node: Node, ind: int) -> None: if "custom" in node.meta: node.meta["custom"].update({"stream": ind}) else: node.meta["custom"] = {"stream": ind} def insert_record_event_after_node(graph: Graph, node: Node, event_ind: int) -> Node: with graph.inserting_after(node): node = graph.call_function( torch.ops.streams.record_event.default, ( event_ind, get_stream_or_current_stream(node), ), ) node.meta["partitioner_tag"] = "must_be_in_backward" return node def insert_wait_event_before_node(graph: Graph, node: Node, event_ind: int) -> Node: with graph.inserting_before(node): node = graph.call_function( torch.ops.streams.wait_event.default, ( event_ind, get_stream_or_current_stream(node), ), ) node.meta["partitioner_tag"] = "must_be_in_backward" return node def populate_stream_timeline( stream_to_timeline: dict[int | None, IndexedDict[Node, float]], graph: Graph, stream_index: int | None, ) -> IndexedDict[Node, float]: if stream_index not in stream_to_timeline: stream_to_timeline[stream_index] = IndexedDict() total_time = 0.0 for node in graph.nodes: # mlazos: not sure if we should include forward here too but don't think it matters if ( node.op == "call_function" and is_bwd_node(node) and get_stream(node) == stream_index ): total_time += get_roofline_estimate(node) stream_to_timeline[stream_index][node] = ( total_time # NB: total time includes the node's runtime ) return stream_to_timeline[stream_index] # NB: we start all estimates at 0, estimating the total runtime of each stream with timestamps at each node # we then try and use these timestamps to estimate when to deallocate tensors used in side streams # See https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html#torch.Tensor.record_stream # for details on the problem being addressed. Rather than using the automatic memory management approach of record_stream # we attempt to find the point which to deallocate based on the estimated timestamps. def handle_synced_deallocation( graph: Graph, stream_to_exec_trace: dict[int | None, IndexedDict[Node, float]], node: Node, last_usage: Node, ) -> None: if not is_bwd_node(node): raise AssertionError( "synced allocations should only be handled on backward nodes" ) if not is_bwd_node(last_usage): raise AssertionError( "synced allocations should only be handled on backward nodes" ) allocating_stream = get_stream(node) side_stream = get_stream(last_usage) if allocating_stream == side_stream: raise AssertionError( "allocating and side stream should be different for synced deallocations" ) if not torch.accelerator.is_available(): # fallback to record_stream in this case with graph.inserting_after(node): rs = graph.call_function( torch.ops.streams.record_stream.default, ( node, get_stream_or_current_stream(last_usage), ), {}, ) rs.meta["partitioner_tag"] = "must_be_in_backward" allocating_stream_trace = populate_stream_timeline( stream_to_exec_trace, graph, allocating_stream ) side_stream_trace = populate_stream_timeline( stream_to_exec_trace, graph, side_stream ) alloc_ptr = node target_side_stream_time = side_stream_trace[last_usage] # linear search from first usage of tensor to a point in time after the side stream has finished while alloc_ptr is not None: alloc_time = allocating_stream_trace[alloc_ptr] if alloc_time >= target_side_stream_time: break elif alloc_time < target_side_stream_time: next_ptr = allocating_stream_trace.next_key(alloc_ptr) if next_ptr is not None: alloc_ptr = next_ptr else: break wait_event = new_event() record_node = insert_record_event_after_node(graph, last_usage, wait_event) with graph.inserting_after(max(alloc_ptr, record_node)): sd = graph.call_function( torch.ops.streams.sync_dealloc.default, (wait_event, get_stream_or_current_stream(alloc_ptr), node), {}, ) sd.meta["partitioner_tag"] = "must_be_in_backward" def insert_sync( graph: Graph, consumer: Node, producer: Node, node_to_wait_event_ind: dict[Node, int], ) -> None: if producer not in node_to_wait_event_ind: node_to_wait_event_ind[producer] = new_event() insert_record_event_after_node( graph, producer, node_to_wait_event_ind[producer] ) insert_wait_event_before_node(graph, consumer, node_to_wait_event_ind[producer]) def assign_backward_streams(gm: torch.fx.GraphModule) -> None: """Assigns backward streams to gradient accumulation nodes""" # NB: iterate in reverse order to more closely match eager # the user node stream will be populated first for node in reversed(list(gm.graph.nodes)): if is_gradient_acc(node): # Accumulation stream selection. Follow the rules from top to bottom to determine the accumulation stream: # 1. Match first stream assignment of the first user with a stream # 2. Match first stream assignment encountered in the args from left to right # This differs from eager in some cases: # Specifically the eager code uses the autograd node to determine the stream, # crucially this does not necessarily correspond to the FX graph node. For example, # in the backward for an add node with a constant we will passthrough and during backward tracing, # no op will be added to the FX graph, so our stream assignment will differ in this case. gradients = _get_flat_args(node, {}) users = list(node.users.keys()) # All gradients will be on same device, they will be coerced if they were not with a .to() node for neighbor in users + gradients: ind = get_stream(neighbor) if ind is not None: set_stream(node, ind) break def insert_backward_syncs(gm: torch.fx.GraphModule) -> None: """Inserts stream syncs for backward nodes if consumer and producer are on different streams""" node_to_wait_event_ind: dict[Node, int] = {} for node in gm.graph.nodes: if node.op == "call_function" and is_bwd_node(node): flat_args = _get_flat_args(node, {}) cur_node_stream = get_stream(node) for arg in flat_args: if arg.op == "call_function" and is_bwd_node(arg): arg_stream = get_stream(arg) if arg_stream != cur_node_stream and get_device(arg).type != "cpu": insert_sync(gm.graph, node, arg, node_to_wait_event_ind) def sync_deallocations(gm: torch.fx.GraphModule) -> None: """Handles https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html#torch.Tensor.record_stream""" # Note: this is only needed if the last usage of a tensor is on a stream other than # the stream the tensor was allocated on # an estimated timestamp from the beginning of graph execution (assuming 0 CPU overhead) # I think this is fine because you should have large tensors if you're using streams # although perhaps I could add a constant 10us per op ahead of the first stream op? # a trace of all the nodes running in a given stream stream_to_exec_trace: dict[int | None, IndexedDict[Node, float]] = {} for node in gm.graph.nodes: if node.op == "call_function" and is_bwd_node(node): allocating_stream = get_stream(node) users = list(node.users.keys()) if not users: continue last_user = max(user for user in users) if last_user.op == "output": continue side_stream = get_stream(last_user) if allocating_stream != side_stream: handle_synced_deallocation( gm.graph, stream_to_exec_trace, node, last_user ) def assign_epilogue_copy_streams(gm: torch.fx.GraphModule) -> None: for epi_copy in gm.graph.find_nodes(op="call_function", target=aten.copy_.default): arg_stream = get_stream(epi_copy.args[1]) copy_stream = get_stream(epi_copy) if arg_stream != copy_stream: set_stream(epi_copy, get_stream_or_current_stream(epi_copy.args[1])) def populate_fw_metadata_with_stream_indices( gm: torch.fx.GraphModule, fw_metadata: "ViewAndMutationMeta" ) -> None: """ Populates fw_metadata.mutated_inp_stream_indices with stream indices from the compiled graph. The forward graph outputs are structured as: (*mutated_inputs, *user_outputs, *intermediate_bases, *saved_tensors, *saved_symints) We extract the stream index for each mutated input from the graph's output node. """ num_mutated_inps = fw_metadata.num_mutated_inp_runtime_indices if num_mutated_inps == 0: fw_metadata.mutated_inp_stream_indices = [] return # Find the output node in the graph output_node = None for node in gm.graph.find_nodes(op="output"): output_node = node break if output_node is None: raise AssertionError( "No output node found in the graph when extracting stream indices" ) # The output node's args[0] is a tuple/list of all outputs output_args = output_node.args[0] # Extract stream indices for the first num_mutated_inps outputs stream_indices = [] for i in range(num_mutated_inps): if i < len(output_args): output_arg = output_args[i] # Get the stream index from the node metadata stream_idx = ( get_stream(output_arg) if isinstance(output_arg, torch.fx.Node) else None ) stream_indices.append(stream_idx) else: stream_indices.append(None) fw_metadata.mutated_inp_stream_indices = stream_indices def _expand_dict_returning_deps( deps: list[Node], visited: set[Node], graph: torch.fx.Graph, sync_node: Node, ) -> list[Node]: """Expand triton_kernel_wrapper_functional dict deps into per-key getitems. ``triton_kernel_wrapper_functional`` returns a flat ``dict[str, Tensor]``. If such a node is threaded through ``control_deps`` as a pass-through value, ``decompose_triton_kernel_wrapper_functional`` later replaces it with a raw Python dict, corrupting every user that expected an FX Node. To avoid this, we insert ``operator.getitem`` nodes *before* the sync for each after-sync getitem user, so that only tensor-valued nodes are passed through ``control_deps``. The dict node may still appear in ``additional_deps`` (ordering-only) which is safe since that tuple is never decomposed. """ expanded: list[Node] = [] for dep in deps: if not ( dep.op == "call_function" and dep.target is torch.ops.higher_order.triton_kernel_wrapper_functional ): expanded.append(dep) continue after_sync_getitems = [ user for user in dep.users if user not in visited and user.op == "call_function" and user.target is operator.getitem ] if not after_sync_getitems: expanded.append(dep) continue with graph.inserting_before(sync_node): for old_gi in after_sync_getitems: key = old_gi.args[1] new_gi = graph.call_function(operator.getitem, args=(dep, key)) new_gi.meta.update(old_gi.meta) if isinstance(new_gi.meta.get("val"), dict): raise AssertionError( f"nested dict in triton kernel output for key {key}" ) visited.add(new_gi) expanded.append(new_gi) old_gi.replace_all_uses_with(new_gi) graph.erase_node(old_gi) return expanded def _wrap_sync_node( gm: torch.fx.GraphModule, sync_node: Node, deps_before_sync: list[Node], visited: set[Node], partition_scoped_deps: set[Node] | None = None, ) -> tuple[Node, list[Node]]: """ Core logic: wrap a single sync node in control_deps. Returns (control_deps_node, passthrough_getitems) where passthrough_getitems are the getitem nodes that thread dependencies through the control_deps node. ``visited`` is the set of nodes at or before the sync node in graph order, used to distinguish pre-sync vs post-sync users. """ from torch._inductor.fx_passes.control_dependencies import ( _create_subgraph_for_node, control_deps, get_subgraph_name, ) graph = gm.graph if partition_scoped_deps is None: partition_scoped_deps = set() sync_is_bwd = is_bwd_node(sync_node) def _user_rewrite_pred(dep: Node) -> Callable[[Node], bool]: if dep in partition_scoped_deps: def pred(user: Node) -> bool: if user in visited: return False if user.op == "output": return sync_is_bwd return is_bwd_node(user) == sync_is_bwd return pred dep_is_bwd = is_bwd_node(dep) def pred(user: Node) -> bool: return user not in visited and (user.op != "output" or dep_is_bwd) return pred # Use dep.users to find deps with uses after the sync — avoids a forward walk. deps_with_uses_after_sync = [] for dep in deps_before_sync: pred = _user_rewrite_pred(dep) if any(pred(user) for user in dep.users): deps_with_uses_after_sync.append(dep) # Expand dict-returning nodes (triton_kernel_wrapper_functional) into # individual getitem outputs. Passing a dict-valued node through # control_deps breaks post-grad passes: reinplace + decomposition replace # the dict node with a Python dict, corrupting any user that expected an # FX Node. Unwrapping here ensures only tensor-valued nodes flow through. deps_with_uses_after_sync = _expand_dict_returning_deps( deps_with_uses_after_sync, visited, graph, sync_node ) # Create subgraph that executes sync and passes through only used dependencies subgraph_module = _create_subgraph_for_node( graph, sync_node, deps_with_uses_after_sync ) subgraph_attr_name = get_subgraph_name(gm, sync_node.name) setattr(gm, subgraph_attr_name, subgraph_module) # Create control_deps call # Note: sync nodes (record_event/wait_event) only take int args, no Node args. with graph.inserting_before(sync_node): get_subgraph = graph.get_attr(subgraph_attr_name) control_deps_node = graph.call_function( control_deps, args=( tuple(deps_before_sync), # additional_deps (all deps for ordering) get_subgraph, # subgraph *deps_with_uses_after_sync, # only pass through deps that are used ), kwargs={}, ) # Propagate partitioner_tag so the partitioner can distinguish # forward vs backward sync control_deps during graph extraction. if "partitioner_tag" in sync_node.meta: control_deps_node.meta["partitioner_tag"] = sync_node.meta["partitioner_tag"] # Mark newly created nodes as visited so subsequent syncs don't # misclassify them as "after the sync" during replacement. visited.add(get_subgraph) visited.add(control_deps_node) # The output is (sync_result, *deps_with_uses_after_sync) # Reuse the val the subgraph output already carries: the min-cut # partitioner's _size_of() raises on any node without `val`, and its escape # hatch covers only get_attr and zero-return OpOverloads, not a schema-less # HOP. The subgraph output is unwrapped when it has no pass-through deps. sg_val = subgraph_module.graph.output_node().meta.get("val") control_deps_node.meta["val"] = sg_val if isinstance(sg_val, tuple) else (sg_val,) # Create getitem nodes only for dependencies that have uses after sync replacements: dict[Node, Node] = {} with graph.inserting_after(control_deps_node): for i, dep in enumerate(deps_with_uses_after_sync): getitem_node = graph.call_function( operator.getitem, args=(control_deps_node, i + 1), # +1 because index 0 is sync result ) getitem_node.meta.update(dep.meta) # The getitem is a passthrough projection of an existing value, not # a node that *produces* an unbacked symbol - the original ``dep`` # (still inside the subgraph) is the binding site. Leaving the # binding on the getitem trips Inductor's run_node assertion that # ``new_unbacked_defs >= renamed_unbacked_bindings``. getitem_node.meta.pop("unbacked_bindings", None) replacements[dep] = getitem_node visited.add(getitem_node) # Replace uses of dependencies that come after sync_node. # Use map_arg to handle nested structures (e.g. output node's list args). for dep, getitem_node in replacements.items(): pred = _user_rewrite_pred(dep) for user in list(dep.users.keys()): if user is control_deps_node: continue if not pred(user): continue def _replace(n: Node) -> Node: return getitem_node if n is dep else n if ( user.op == "call_function" and user.target is torch.ops.aten.copy_.default and user.args[0] is dep and dep in partition_scoped_deps ): # AOT functionalization leaves keep_input_mutations epilogue # copy_ as the only writer of a partition-scoped placeholder. # Preserve the destination identity used by mutation bookkeeping. user.args = (dep, *map_arg(user.args[1:], _replace)) else: user.args = map_arg(user.args, _replace) user.kwargs = map_arg(user.kwargs, _replace) # Remove original sync node sync_node.replace_all_uses_with(control_deps_node) graph.erase_node(sync_node) return control_deps_node, list(replacements.values()) def _collect_sync_forward_deps( graph: torch.fx.Graph, ) -> tuple[dict[Node, OrderedSet[Node]], dict[Node, OrderedSet[Node]]]: """Collect values that must cross sync boundaries. ``live_inputs[partition][stream]`` tracks values whose passthrough can anchor later work in that partition on that stream. Demand labels propagate through producers while each producer also seeds its own execution stream. This is a single reverse liveness traversal; the forward pass still handles sync/event state introduced while rewriting the graph. """ stream_sync_deps: dict[Node, OrderedSet[Node]] = {} full_barrier_deps: dict[Node, OrderedSet[Node]] = {} live_inputs: dict[Partition, dict[int, OrderedSet[Node]]] = { FORWARD: {}, BACKWARD: {}, } locations: dict[Node, OrderedSet[tuple[Partition, int]]] = {} full_barriers = ( torch.ops.streams.synchronize_device.default, torch.ops.streams.synchronize_event.default, torch.ops.streams.synchronize_stream.default, ) def _add_live(node: Node, partition: Partition, stream: int) -> None: stream_inputs = live_inputs[partition].get(stream) if stream_inputs is None: stream_inputs = live_inputs[partition][stream] = OrderedSet() stream_inputs.add(node) node_locations = locations.get(node) if node_locations is None: node_locations = locations[node] = OrderedSet() node_locations.add((partition, stream)) for node in reversed(graph.nodes): if node.op == "call_function" and node.target in full_barriers: deps = OrderedSet( dep for stream_inputs in live_inputs[is_bwd_node(node)].values() for dep in stream_inputs ) if deps: full_barrier_deps[node] = deps continue if ( node.op == "call_function" and node.target is torch.ops.streams.wait_stream.default ): waiting_stream: int = node.args[0] # type: ignore[assignment] deps = live_inputs[is_bwd_node(node)].get(waiting_stream) if deps: stream_sync_deps[node] = deps.copy() continue if ( node.op == "call_function" and node.target is torch.ops.streams.wait_event.default ): waiting_stream: int = node.args[1] # type: ignore[assignment] deps = live_inputs[is_bwd_node(node)].get(waiting_stream, ()) graph_inputs = OrderedSet(dep for dep in deps if dep.op == "placeholder") if graph_inputs: stream_sync_deps[node] = graph_inputs continue if node.op != "call_function" or node.target in _SYNC_OPS: continue inherited_locations = locations.pop(node, _EMPTY_LOCATIONS) for partition, stream in inherited_locations: live_inputs[partition][stream].discard(node) execution_stream = get_stream(node) if execution_stream is None: execution_stream = 0 target_locations = OrderedSet( (*inherited_locations, (is_bwd_node(node), execution_stream)) ) def _collect(n: Node) -> Node: if "val" in n.meta: for partition, stream in target_locations: _add_live(n, partition, stream) return n map_arg(node.args, _collect) map_arg(node.kwargs, _collect) return stream_sync_deps, full_barrier_deps def _collect_full_barrier_deps( stream_to_nodes: dict[int | None, list[Node]], stream_sync_deps: dict[int | None, list[Node]], bwd: bool = False, ) -> list[Node]: """Collect all still-live cross-stream deps for a full CPU barrier (synchronize_stream/device/event). These ops block the CPU, so a consumer is only ordered after them by an explicit control_deps edge (nothing on the consumer's stream waits for a CPU barrier). Includes live compute nodes (stream_to_nodes) plus the ctrl/ passthrough nodes accumulated by prior sync ops that already cleared stream_to_nodes (stream_sync_deps), deduped by identity preserving order. ``bwd`` marks a backward barrier: its accumulated sync deps are filtered to backward nodes only, so it never depends on a forward control_deps node. That edge would cross the fwd/bwd boundary onto a node that is neither saveable nor recomputable, making the min-cut partitioner infeasible. Forward compute deps (stream_to_nodes) are kept -- backward legitimately depends on forward activations. """ all_deps = [n for nodes in stream_to_nodes.values() for n in nodes] seen = set(all_deps) for deps in stream_sync_deps.values(): for dep in deps: if dep in seen or (bwd and not is_bwd_node(dep)): continue all_deps.append(dep) seen.add(dep) return all_deps def wrap_all_sync_nodes_with_control_deps(gm: torch.fx.GraphModule) -> None: """ Single-pass wrap of all sync nodes in control_deps. Iterates through the graph once, accumulating per-stream node lists. When a sync node is encountered, it is wrapped using the accumulated deps for that stream, then the deps are reset to the control_deps node (maintaining the ordering chain for subsequent syncs on the same stream). """ graph = gm.graph if len(graph.nodes) == 0: raise RuntimeError("Expected a non-empty graph") stream_sync_forward_deps, full_barrier_forward_deps = _collect_sync_forward_deps( graph ) stream_to_nodes: dict[int | None, list[Node]] = {} # Maps event_index to its latest ordering anchor: either the control_deps that # wrapped record_event or a raw record_event with no dependencies. event_to_anchor: dict[int, Node] = {} # Maps event_index -> getitem nodes threaded through record_event's control_deps, # so synchronize_event can thread them through to subsequent ops. event_to_passthrough: dict[int, list[Node]] = {} # Maps event_index -> stream that the event was recorded on, # so synchronize_event can infer its stream. event_to_stream: dict[int, int | None] = {} # Tracks the most recent control_deps node and its passthrough # getitems per stream. synchronize_stream/device consults this to # depend on prior sync ops (ctrl_node for ordering) and thread # cross-stream data through (passthrough for data edges), even # after stream_to_nodes has been cleared by those sync ops. stream_sync_deps: dict[int | None, list[Node]] = {} visited: set[Node] = set() found_sync = False # Walk the node linked-list manually so we can mutate the graph # (wrapping sync nodes inserts/erases nodes) without losing our place. node = next(iter(graph.nodes)) while node.op != "root": next_node = node.next visited.add(node) if node.op == "call_function": if node.target in _SYNC_OPS: # synchronize_device and synchronize_stream block the CPU, # so all subsequent kernel launches are host-ordered after # them. Treat both as full barriers across all streams. if node.target in ( torch.ops.streams.synchronize_device.default, torch.ops.streams.synchronize_stream.default, ): all_stream_deps = _collect_full_barrier_deps( stream_to_nodes, stream_sync_deps, bwd=is_bwd_node(node) ) existing_deps = set(all_stream_deps) for dep in full_barrier_forward_deps.get(node, ()): if dep not in existing_deps: all_stream_deps.append(dep) existing_deps.add(dep) if all_stream_deps: found_sync = True ctrl_node_sync, passthrough_sync = _wrap_sync_node( gm, node, all_stream_deps, visited ) else: ctrl_node_sync = None passthrough_sync: list[Node] = [] stream_to_nodes.clear() stream_sync_deps.clear() sync_anchor = ctrl_node_sync if ctrl_node_sync is not None else node stream_sync_deps[None] = [sync_anchor, *passthrough_sync] while ( getattr(next_node, "_erased", False) and next_node.op != "root" ): next_node = next_node.next node = next_node continue # wait_stream(waiting, waited_on) creates a dependency on # the waited_on stream: all prior work there must complete # before the waiting stream proceeds. if node.target is torch.ops.streams.wait_stream.default: waiting_stream: int = node.args[0] # type: ignore[assignment] waited_on_stream: int = node.args[1] # type: ignore[assignment] deps_before_sync = list(stream_to_nodes.get(waited_on_stream, ())) if waiting_stream != waited_on_stream: deps_before_sync.extend(stream_to_nodes.get(waiting_stream, ())) if None in stream_to_nodes and waited_on_stream is not None: deps_before_sync.extend(stream_to_nodes[None]) participating_streams = dict.fromkeys( (waited_on_stream, waiting_stream, None) ) prior_syncs = list( dict.fromkeys( dep for stream in participating_streams for dep in stream_sync_deps.get(stream, ()) ) ) if is_bwd_node(node): prior_syncs = [n for n in prior_syncs if is_bwd_node(n)] deps_before_sync.extend(prior_syncs) # Also include inputs of subsequent waiting-stream nodes # so they get threaded through control_deps, creating a # forward ordering constraint (waiting-stream work must # come after wait_stream). existing_deps = set(deps_before_sync) for dep in stream_sync_forward_deps.get(node, ()): if dep not in existing_deps: deps_before_sync.append(dep) existing_deps.add(dep) if deps_before_sync: found_sync = True ctrl_node_ws, passthrough_ws = _wrap_sync_node( gm, node, deps_before_sync, visited ) else: ctrl_node_ws = None passthrough_ws: list[Node] = [] wait_anchor = ctrl_node_ws if ctrl_node_ws is not None else node # Key under both streams: the wait executes on waiting_stream # and subsumes the waited-on work used to anchor it. stream_sync_deps[waiting_stream] = [ wait_anchor, *passthrough_ws, ] stream_sync_deps[waited_on_stream] = [ wait_anchor, *passthrough_ws, ] stream_to_nodes[waited_on_stream] = [] if None in stream_to_nodes: stream_to_nodes[None] = [] while ( getattr(next_node, "_erased", False) and next_node.op != "root" ): next_node = next_node.next node = next_node continue event_index: int = node.args[0] # type: ignore[assignment] # synchronize_event blocks the CPU thread, so it acts # as a barrier across all streams. Collect deps from every # stream and reset them all afterward. If the event was # recorded externally, thread the graph inputs through so # that any post-sync uses depend on the synchronize. partition_scoped_deps: list[Node] = [] if node.target is torch.ops.streams.synchronize_event.default: sync_stream: int | None = event_to_stream.get(event_index) # Like synchronize_stream/device, this is a CPU-side barrier, # so fold in still-live cross-stream data accumulated by prior # sync ops (which cleared stream_to_nodes). all_stream_deps = _collect_full_barrier_deps( stream_to_nodes, stream_sync_deps, bwd=is_bwd_node(node) ) deps_before_sync = [ *all_stream_deps, *full_barrier_forward_deps.get(node, ()), ] else: sync_stream = node.args[1] # type: ignore[assignment] deps_before_sync = list(stream_to_nodes.get(sync_stream, ())) partition_scoped_deps = ( list(stream_sync_forward_deps.get(node, ())) if node.target is torch.ops.streams.wait_event.default and event_index not in event_to_anchor else [] ) deps_before_sync.extend(partition_scoped_deps) # Nodes without explicit stream annotation (custom.stream=None) # run on the current/default stream. Include them when the sync # op references a stream, since the unannotated nodes are # implicitly on that stream. if None in stream_to_nodes and sync_stream is not None: deps_before_sync.extend(stream_to_nodes[None]) # Order this sync op after the prior sync on its own stream # (and after any full barrier, keyed under None). Without this, # a record/wait whose stream_to_nodes was reset by an earlier # sync would have no deps, stay unwrapped, and float as a bare # side-effectful node -- capturing nothing, the same failure as # an unanchored synchronize_stream. prior_syncs = [ *stream_sync_deps.get(sync_stream, ()), *stream_sync_deps.get(None, ()), ] # A backward sync must anchor only on backward sync nodes: # depending on a forward control_deps node crosses the fwd/bwd # boundary onto a node that is neither saveable nor # recomputable, which makes the min-cut partitioner infeasible. if is_bwd_node(node): prior_syncs = [n for n in prior_syncs if is_bwd_node(n)] deps_before_sync.extend(prior_syncs) # For wait_event and synchronize_event, depend on the matching # record_event's control_deps node (ordering) and thread its # passthrough getitems through (data edge), so consumers of # recorded values are rewired through the wait/synchronize # and the scheduler sees the dependency. if node.target in ( torch.ops.streams.wait_event.default, torch.ops.streams.synchronize_event.default, ): if event_index in event_to_anchor: deps_before_sync = [ event_to_anchor[event_index], *deps_before_sync, ] if event_index in event_to_passthrough: deps_before_sync = [ *deps_before_sync, *event_to_passthrough[event_index], ] elif ( node.target is torch.ops.streams.record_event.default and event_index in event_to_anchor ): # Re-recording replaces an event's state, but the host-side # record operations must retain program order even when the # new record is on an otherwise-empty stream. deps_before_sync.insert(0, event_to_anchor[event_index]) # Deduplicate (preserving order): the folded stream_sync_deps and # the event's own anchor/passthrough can overlap. deps_before_sync = list(OrderedSet(deps_before_sync)) if deps_before_sync: found_sync = True ctrl_node, passthrough = _wrap_sync_node( gm, node, deps_before_sync, visited, set(partition_scoped_deps) if node.target is torch.ops.streams.wait_event.default else None, ) else: ctrl_node = None passthrough: list[torch.fx.Node] = [] if node.target is torch.ops.streams.record_event.default: event_to_stream[event_index] = sync_stream event_to_anchor[event_index] = ( ctrl_node if ctrl_node is not None else node ) event_to_passthrough[event_index] = passthrough sync_anchor = ctrl_node if ctrl_node is not None else node stream_sync_deps[sync_stream] = [sync_anchor, *passthrough] # Reset: ops between this sync and the next will accumulate # fresh. Ordering with prior ops is already enforced because # their uses were rewired through getitems from control_deps. # synchronize_event is a full barrier, so its own passthrough # subsumes all prior per-stream data: clear then re-seed. Key the # reseed under None (like synchronize_stream/device) so a later # record/wait on any stream chains after it. if node.target is torch.ops.streams.synchronize_event.default: stream_to_nodes.clear() stream_sync_deps.clear() stream_sync_deps[None] = [sync_anchor, *passthrough] else: stream_to_nodes[sync_stream] = [] if None in stream_to_nodes: stream_to_nodes[None] = [] elif "val" in node.meta: stream = get_stream(node) stream_to_nodes.setdefault(stream, []).append(node) # _wrap_sync_node may erase nodes after the sync (e.g. # _expand_dict_returning_deps erases getitems). Skip erased nodes # so we don't re-process a dead node into stream_to_nodes. while getattr(next_node, "_erased", False) and next_node.op != "root": next_node = next_node.next node = next_node if found_sync: gm.recompile()