1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
|
#pragma once
#include <torch/csrc/WindowsTorchApiMacro.h>
#include <torch/csrc/jit/ir/alias_analysis.h>
#include <torch/csrc/jit/ir/ir.h>
namespace torch {
namespace jit {
// Utilities for dealing with nodes that contain subgraphs.
//
// They handle the complexity of editing inputs/outputs as you merge nodes in
// and out of subgraphs.
namespace SubgraphUtils {
// Create a new subgraph node that contains only `n`. The new subgraph will have
// `subgraphKind` as its type.
//
// `n` is destroyed.
//
// Returns the new subgraph node.
// An optional argument 'vmap' could be used to retrieve value mappings
// Values will be mapped to their new subgraph values
TORCH_API Node* createSingletonSubgraph(Node* n, Symbol subgraphKind);
TORCH_API Node* createSingletonSubgraph(
Node* n,
Symbol subgraphKind,
std::unordered_map<Value*, Value*>& vmap);
// Creates a new subgraph that only contains `n`, amd udpates the new outputs
// of the subgraph to have the aliasing properties of the original `n` outputs
TORCH_API Node* createSingletonSubgraphAndUpdateAliasing(
Node* to_merge,
Symbol subgraphKind,
AliasDb& db);
// Merge a node into a subgraph node. If `toMerge` is also a subgraph, the
// subgraphs are merged.
// If `destroyNode` is true `toMerge` is destroyed.
// An optional argument 'vmap' could be used to retrieve value mappings.
// Values will be mapped to their new subgraph values
TORCH_API void mergeNodeIntoSubgraph(
Node* toMerge,
Node* subgraphNode,
bool destroyNode = true);
TORCH_API void mergeNodeIntoSubgraph(
Node* toMerge,
Node* subgraphNode,
std::unordered_map<Value*, Value*>& vmap,
bool destroyNode = true);
// Merges a node into a subgraph node, and updates the new outputs of the
// subgraph to have the aliasing properties of the corresponding `to_merge`
// outputs
TORCH_API void mergeNodeIntoSubgraphAndUpdateAliasing(
Node* to_merge,
Node* subgraphNode,
AliasDb& db);
// Move nodes from a subgraph node to the outer graph.
// `subgraphNode` is destroyed.
// An optional argument 'vmap' could be used to retrieve value mappings.
TORCH_API void unmergeSubgraph(Node* subgraphNode);
TORCH_API void unmergeSubgraph(
Node* subgraphNode,
std::unordered_map<Value*, Value*>& vmap);
// Convenience function
std::shared_ptr<Graph> getSubgraph(Node* n);
} // namespace SubgraphUtils
} // namespace jit
} // namespace torch
|
