|
8 | 8 |
|
9 | 9 | import uuid |
10 | 10 | from collections import defaultdict |
| 11 | +from contextlib import contextmanager |
11 | 12 | from dataclasses import dataclass |
12 | 13 | from enum import auto, Enum |
13 | | -from typing import Dict, Optional, Tuple, Union |
| 14 | +from typing import Dict, Iterator, List, Optional, Tuple, Union |
14 | 15 |
|
15 | 16 | import torch |
16 | 17 | from torch.fx.node import Node |
@@ -73,10 +74,72 @@ def __init__(self): |
73 | 74 | self.tid_managers: Dict[IdSpace, IdManager] = defaultdict(IdManager) |
74 | 75 | self.vid_managers: Dict[IdSpace, IdManager] = defaultdict(IdManager) |
75 | 76 | self.name_to_slot: Dict[str, Slot] = {} |
| 77 | + # Stack of active temp-slot scopes (see ``tmp_scope``). Temp tids/vids |
| 78 | + # allocated via make_tmp_slot()/make_tmp_value_slot() are registered on |
| 79 | + # the innermost scope and their ids returned for reuse on scope exit. |
| 80 | + self._tmp_scopes: List[List[Slot]] = [] |
| 81 | + |
| 82 | + @contextmanager |
| 83 | + def tmp_scope(self) -> Iterator[None]: |
| 84 | + """Scope temporary slot allocations so their ids can be reused. |
| 85 | +
|
| 86 | + Temp tids/vids allocated via :meth:`make_tmp_slot` / |
| 87 | + :meth:`make_tmp_value_slot` inside this context are returned to their |
| 88 | + id pools when the context exits, so later allocations (temp or node) |
| 89 | + can reuse them. Allocating a temp slot outside any ``tmp_scope`` raises |
| 90 | + ``RuntimeError``. |
| 91 | +
|
| 92 | + Scopes may be nested; each allocation is tied to the innermost scope. |
| 93 | + The Slot objects stay in ``name_to_slot`` (mirroring node-slot reclaim |
| 94 | + via ``return_id``) so serialization still sees every distinct slot. |
| 95 | +
|
| 96 | + Invariant: the slots yielded here are dead after the context exits. |
| 97 | + Never ``set_slot`` a temp slot as a node's persistent output because its id is |
| 98 | + reclaimed on scope exit and would be reused (and coalesced) by a later |
| 99 | + node while still live. Node outputs must come from ``make_or_get_slot``. |
| 100 | + """ |
| 101 | + self._tmp_scopes.append([]) |
| 102 | + try: |
| 103 | + yield |
| 104 | + finally: |
| 105 | + scope = self._tmp_scopes.pop() |
| 106 | + for slot in scope: |
| 107 | + if slot.id_type == IdType.Tensor: |
| 108 | + self.tid_managers[slot.id_space].return_id(slot.idx) |
| 109 | + else: |
| 110 | + self.vid_managers[slot.id_space].return_id(slot.idx) |
| 111 | + |
| 112 | + def _new_tmp_slot(self, id_type: IdType, prefix: str) -> Tuple[str, Slot]: |
| 113 | + if not self._tmp_scopes: |
| 114 | + raise RuntimeError( |
| 115 | + f"{prefix}() must be called within a SlotManager.tmp_scope() " |
| 116 | + "context so temporary ids can be reclaimed and reused." |
| 117 | + ) |
| 118 | + name = f"{prefix}_{uuid.uuid4().hex}" |
| 119 | + id_space = IdSpace.Temp |
| 120 | + manager = ( |
| 121 | + self.tid_managers[id_space] |
| 122 | + if id_type == IdType.Tensor |
| 123 | + else self.vid_managers[id_space] |
| 124 | + ) |
| 125 | + idx = manager.get_id() |
| 126 | + slot = Slot(id_type=id_type, id_space=id_space, idx=idx) |
| 127 | + self.name_to_slot[name] = slot |
| 128 | + self._tmp_scopes[-1].append(slot) |
| 129 | + return name, slot |
76 | 130 |
|
77 | 131 | def set_slot(self, node_or_name: Union[Node, str], slot: Slot): |
78 | 132 | if isinstance(node_or_name, Node): |
79 | 133 | node_or_name = node_or_name.name |
| 134 | + # A slot still tracked by an active tmp_scope has its id reclaimed when the |
| 135 | + # scope exits, so it must never be bound as a node's persistent output (a |
| 136 | + # later node would read it as dead). Node outputs must come from |
| 137 | + # make_or_get_slot(). See SlotManager.tmp_scope(). |
| 138 | + assert not any(slot in scope for scope in self._tmp_scopes), ( |
| 139 | + f"Cannot bind temporary slot {slot} as the output of {node_or_name}; " |
| 140 | + f"its id is reclaimed on tmp_scope exit. Use make_or_get_slot() for " |
| 141 | + f"node outputs." |
| 142 | + ) |
80 | 143 | # Allow setting a slot to the same value (e.g., for in-place ops like SLICE_UPDATE) |
81 | 144 | existing = self.name_to_slot.get(node_or_name) |
82 | 145 | if existing is not None: |
@@ -129,23 +192,11 @@ def make_constant_slot(self, name: str) -> Slot: |
129 | 192 | return slot |
130 | 193 |
|
131 | 194 | def make_tmp_slot(self) -> Tuple[str, Slot]: |
132 | | - name = f"tmp_{uuid.uuid4().hex}" |
133 | | - id_space = IdSpace.Temp |
134 | | - manager = self.tid_managers[id_space] |
135 | | - idx = manager.get_id() |
136 | | - slot = Slot(id_type=IdType.Tensor, id_space=id_space, idx=idx) |
137 | | - self.name_to_slot[name] = slot |
138 | | - return name, slot |
| 195 | + return self._new_tmp_slot(IdType.Tensor, "tmp") |
139 | 196 |
|
140 | 197 | def make_tmp_value_slot(self) -> Tuple[str, Slot]: |
141 | 198 | """Create a temporary SymInt slot and register it.""" |
142 | | - name = f"tmp_val_{uuid.uuid4().hex}" |
143 | | - id_space = IdSpace.Temp |
144 | | - manager = self.vid_managers[id_space] |
145 | | - idx = manager.get_id() |
146 | | - slot = Slot(id_type=IdType.SymInt, id_space=id_space, idx=idx) |
147 | | - self.name_to_slot[name] = slot |
148 | | - return name, slot |
| 199 | + return self._new_tmp_slot(IdType.SymInt, "tmp_val") |
149 | 200 |
|
150 | 201 | def make_or_get_slots( |
151 | 202 | self, node: Node, id_space: IdSpace = IdSpace.Temp |
|
0 commit comments