| 273 | * being executed. |
| 274 | */ |
| 275 | export function getNodeLiveUntilMap(orderedNodes: Node[]): Map<string, Node[]> { |
| 276 | const nodeNameToOrder = new Map<string, number>( |
| 277 | orderedNodes.map((node, order) => [node.name, order])); |
| 278 | |
| 279 | const INF_LIFE = Number.MAX_SAFE_INTEGER; |
| 280 | // Make control flow nodes (and consequently their direct parents) |
| 281 | // live forever since they're tricky to track correctly. |
| 282 | const selfLifespans = orderedNodes.map( |
| 283 | (node, nodeOrder) => isControlFlow(node) ? INF_LIFE : nodeOrder); |
| 284 | const getSelfLifeSpan = (node: Node) => { |
| 285 | const selfLife = selfLifespans[nodeNameToOrder.get(node.name)!]; |
| 286 | if (selfLife == null) { |
| 287 | // If nodeToOrder does not contain the node, it is unused or |
| 288 | // unreachable in graph. |
| 289 | return -1; |
| 290 | } |
| 291 | return selfLife; |
| 292 | }; |
| 293 | |
| 294 | // `liveUntil[i]` points to the last node in the `orderedNodes` array that |
| 295 | // may depend on tensors from node `i`. It indicates that all the |
| 296 | // intermediate tensors from `orderedNodes[i]` should be disposed after |
| 297 | // `orderedNodes[liveUntil[i]]` is executed. |
| 298 | // A node lives long enough to pass on its tensors to its children. |
| 299 | // It lives until at least `max(node's position, children's positions)`. |
| 300 | const liveUntilOrders = orderedNodes.map((node, nodeOrder) => { |
| 301 | return node.children.map(getSelfLifeSpan) |
| 302 | .reduce((a, b) => Math.max(a, b), selfLifespans[nodeOrder]); |
| 303 | }); |
| 304 | |
| 305 | // liveUntilMap: |
| 306 | // - Key: Name of a node `x` |
| 307 | // - Values: All nodes whose intermediate tensors should be disposed |
| 308 | // after `x` is executed. |
| 309 | const liveUntilMap = new Map<string, Node[]>(); |
| 310 | for (let nodeOrder = 0; nodeOrder < orderedNodes.length; ++nodeOrder) { |
| 311 | const liveUntilOrder = liveUntilOrders[nodeOrder]; |
| 312 | if (liveUntilOrder === INF_LIFE) { |
| 313 | continue; |
| 314 | } |
| 315 | const node = orderedNodes[nodeOrder]; |
| 316 | const liveUntilNode = orderedNodes[liveUntilOrder]; |
| 317 | if (!liveUntilMap.has(liveUntilNode.name)) { |
| 318 | liveUntilMap.set(liveUntilNode.name, []); |
| 319 | } |
| 320 | liveUntilMap.get(liveUntilNode.name)!.push(node); |
| 321 | } |
| 322 | return liveUntilMap; |
| 323 | } |
| 324 | |
| 325 | const CONTROL_FLOW_OPS = new Set([ |
| 326 | 'Switch', 'Merge', 'Enter', 'Exit', 'NextIteration', 'StatelessIf', |