/usr/local/lib64/python3.6/site-packages/torch/fx/experimental
Edit: /usr/local/lib64/python3.6/site-packages/torch/fx/experimental/const_fold.py (13236B)
import operator
from typing import Dict, Set, List, Optional, Union
import torch.fx
import torch.fx.experimental.fx_acc.acc_utils as acc_utils
from torch.fx.passes.split_module import split_module
def _make_tuple(x):
"""
Helper to convert x into a one item tuple if it's not a tuple already.
"""
return x if isinstance(x, tuple) else (x,)
class FoldedGraphModule(torch.fx.GraphModule):
"""
FoldedGraphModule is a GraphModule which also contains another
`const_subgraph_module` representing a subgraph which has all const attr
inputs and which can be run once before running the main standard
`graph`. The `const_output_names` are the ordered list names of attrs which
represent what each respective output from the const_subgraph should be set
on which attrs.
"""
def __init__(
self,
root: torch.nn.Module,
graph: torch.fx.Graph,
const_subgraph: Optional[torch.fx.Graph] = None,
const_output_names: Optional[List[str]] = None,
):
super().__init__(root, graph)
self.const_subgraph_module = (
None
if const_subgraph is None
else torch.fx.GraphModule(root, const_subgraph)
)
self.const_output_names = const_output_names
self.has_folding_been_run = False
def __call__(self, *args, **kwargs):
if not self.has_folding_been_run:
self.run_folding()
return super().__call__(*args)
def run_folding(self):
# If there's no const subgraph module or attr output names to use, return
# early as there is no const folding to perform.
if self.const_subgraph_module is None or self.const_output_names is None:
return
assert not self.has_folding_been_run
self.has_folding_been_run = True
# Actually run const folding subgraph. We _make_tuple here because
# single attr const fold subgraphs output a single Tensor while
# multiple outputs are returned as Tuple[Tensor,].
folded_attrs = _make_tuple(self.const_subgraph_module())
# Look for output node from const folding subgraph and set attrs on the
# module with the results.
for i in range(len(folded_attrs)):
setattr(
self, self.const_output_names[i], torch.nn.Parameter(folded_attrs[i])
)
def split_const_subgraphs(
module: Union[torch.nn.Module, torch.fx.GraphModule]
) -> FoldedGraphModule:
"""
Looks through `module` for any nodes that have all constant attribute inputs
and separates them out into their own constant subgraph, and returns a
FoldedGraphModule which runs that constant subgraph on the first run to set
attributes on the module prior to running the non-constant portion of the
graph.
"""
if not isinstance(module, torch.fx.GraphModule):
mod_traced = torch.fx.symbolic_trace(module)
else:
mod_traced = module
# Build up a list of const_nodes, defined as nodes that are themselves
# get_attrs, or have all get_attr or other constant node inputs.
const_nodes: Set[torch.fx.Node] = set()
found_const_folding = False
for node in mod_traced.graph.nodes:
# Skip over placeholders/outputs because they can't be const folded and
# we don't want to add tags to them.
if node.op in {"placeholder", "output"}:
continue
# If the node itself is constant, or all of its inputs are constant,
# then tag it as constant.
if node.op == "get_attr" or set(node.all_input_nodes).issubset(const_nodes):
const_nodes.add(node)
if node.op != "get_attr":
found_const_folding = True
# If we did not find any const folding then return early without a const fold subgraph.
if not found_const_folding:
return FoldedGraphModule(mod_traced, mod_traced.graph)
# Partition the module into two: submod_0 for constant folding subgraph, and
# submod_1 for the rest.
def mod_partition(node: torch.fx.Node):
return 0 if node in const_nodes else 1
split = split_module(mod_traced, module, mod_partition)
# Later we are creating get_attr nodes from main module get_attr nodes and we use the
# same node.target. If there's a get_attr that refers to an attributes in a module and
# this module is not included in submod_1 then we would have a trouble when creating
# get_attr ndoes because it will try to find the module that owns the attribute first.
# Setting owning_module here makes the owning_module of submod_1.graph to None then
# the check when creating get_attr nodes will get skipped.
split.submod_1.graph.owning_module = mod_traced
# The module that a call_module node refers to gets copied to submodules during split.
# The path to the module also gets inlined, i.e. mod.a.b -> mod_a_b. Here we need to
# attach inlined modules to `mod_traced` as it's the owning module now.
for node in split.submod_1.graph.nodes:
if node.op == "call_module":
setattr(mod_traced, node.target, getattr(split.submod_1, node.target))
# Gather all names that are output from the const folding subgraph, which we
# will need to set dummy params on the module.
const_output_names: List[str] = []
for node in split.submod_0.graph.nodes:
if node.op == "output":
# Note: we _make_tuple here because the output Node either contains
# a single output Node, or Tuple[Node], so this simplifies things.
const_output_names = [o.name for o in _make_tuple(node.args[0])]
break
# Make sure the attr name we want to use is uniquely named in the module.
for i in range(len(const_output_names)):
# Add a suffix to make it easier to tell these were the result of const folding.
name = const_output_names[i] + "__CF"
# Get an unique name for the attr.
name = acc_utils.get_unique_attr_name_in_module(mod_traced, name)
const_output_names[i] = name
# Now track the const_output_names to what name is used in the parent graph
# from the split via call_function getitem, to see what order it is passed
# into the non-const subgraph submod_1. First look to the parent module
# containing/calling into the const/non-const submodules to determine what
# the inputs are to each. Note if submod_0 had a single output then there is
# no getitem, and we can simply use the output from the call to submoid_0.
call_submod_0_args, call_submod_1_args = None, None
orig_ph_targets: List[str] = []
for node in split.graph.nodes:
if node.op == "placeholder":
orig_ph_targets.append(node.target)
if node.op == "call_module":
if node.target == "submod_0":
call_submod_0_args = node.args
continue
elif node.target == "submod_1":
call_submod_1_args = node.args
continue
assert call_submod_0_args is not None and call_submod_1_args is not None
# Look through the args for the call into submod_1, and find the args that
# come from submod_0. Also look for get_attrs fed directly from the parent
# split into submod_1, i.e. those attrs that are not constant folded.
submod_1_input_idx_to_folded_attr_name: Dict[int, str] = {}
submod_1_input_idx_to_unfolded_attr_name: Dict[int, str] = {}
for i, node in enumerate(call_submod_1_args):
const_output_name = None
# If we only had a single output from submod_0 then we simply look for
# the call_module into it.
if len(const_output_names) == 1:
if node.op == "call_module" and node.target == "submod_0":
const_output_name = const_output_names[0]
# Else we had multiple outputs from submod_0, so we need to look for all
# getitems from the call to it.
else:
if (
node.op == "call_function"
and node.target == operator.__getitem__
and node.args[0].target == "submod_0"
):
const_output_name = const_output_names[node.args[1]]
# Now map from the index of the constant into calling submod_1 and map
# to the constant output name, which we use for swapping in getattrs
# instead of placeholders in submod_1.
if const_output_name is not None:
submod_1_input_idx_to_folded_attr_name[i] = const_output_name
elif node.op == "get_attr":
submod_1_input_idx_to_unfolded_attr_name[i] = node.target
assert len(submod_1_input_idx_to_folded_attr_name) == len(const_output_names)
# Now we have a mapping from const output names to the index they are passed
# into submod_1, so swap in getattrs for placeholders.
ph_idx = 0
for node in split.submod_1.graph.nodes:
if node.op != "placeholder":
continue
is_folded_attr = ph_idx in submod_1_input_idx_to_folded_attr_name
is_unfolded_attr = ph_idx in submod_1_input_idx_to_unfolded_attr_name
if not is_folded_attr and not is_unfolded_attr:
ph_idx += 1
continue
const_output_name = (
submod_1_input_idx_to_folded_attr_name[ph_idx]
if is_folded_attr
else submod_1_input_idx_to_unfolded_attr_name[ph_idx]
)
if is_folded_attr:
assert not hasattr(mod_traced, const_output_name)
# Use a dummy param, which will be overwritten when we run const folding.
setattr(
mod_traced,
const_output_name,
torch.nn.Parameter(torch.randn(1)),
)
with split.submod_1.graph.inserting_before(node):
new_node = split.submod_1.graph.get_attr(const_output_name)
new_node.meta = node.meta.copy()
node.replace_all_uses_with(new_node)
split.submod_1.graph.erase_node(node)
ph_idx += 1
# We may need to reorder placeholders to ensure they have the same order as
# they do in the original split.
ph_idx = 0
node = next(iter(split.submod_1.graph.nodes))
while node.op != "root":
if node.op != "placeholder":
node = node.next
continue
curr_orig_ph_target = orig_ph_targets[ph_idx]
ph_idx += 1
# If this ph is in the correct position, nothing to do.
if curr_orig_ph_target == node.target:
node = node.next
continue
# This ph is not in the correct order, so search the rest of the graph
# for the ph we expected and prepend it before the current ph.
later_node = node.next
while later_node.op != "root":
if (
later_node.op == "placeholder"
and curr_orig_ph_target == later_node.target
):
break
later_node = later_node.next
# The placeholder is in split.graph but not in split.submod_1.graph.
# In this case we add the placeholder to submod_1.graph.
if later_node.op == "root":
with split.submod_1.graph.inserting_before(node):
split.submod_1.graph.placeholder(curr_orig_ph_target)
else:
node.prepend(later_node)
# Note we do not increment node here, as it still may be in the wrong
# place (we just prepended the ph that should have come before it).
# There might be more placeholders left in orig_ph_targets.
last_placeholder = next(n for n in split.submod_1.graph.nodes if n.target == orig_ph_targets[ph_idx - 1])
while ph_idx < len(orig_ph_targets):
with split.submod_1.graph.inserting_after(last_placeholder):
split.submod_1.graph.placeholder(orig_ph_targets[ph_idx])
ph_idx += 1
# split_module currently does not use get_attrs for attrs. Instead it passes
# them in as args from the parent module, which used get_attrs. Here we set
# them as get_attrs inside submod_0, allowing for running folding without
# somehow a priori knowing the attrs that should be passed as args. We can
# unconditionally do this for all placeholders because we know all
# placeholders to submod_0 must be constants accessible via get_attr.
# Note that here we set the split.submod_0.graph into a new root_submod_0 with split
# as the root module, because we are fetching attributes directly from the root
# module, instead of fetching them from split.submod_0.
root_submod_0 = torch.fx.GraphModule(split, split.submod_0.graph)
for node in root_submod_0.graph.nodes:
if node.op != "placeholder":
continue
in_node = next(n for n in call_submod_0_args if n.name == node.target)
assert in_node.op == "get_attr"
with root_submod_0.graph.inserting_before(node):
node.replace_all_uses_with(root_submod_0.graph.get_attr(in_node.target))
root_submod_0.graph.erase_node(node)
return FoldedGraphModule(
mod_traced, split.submod_1.graph, root_submod_0.graph, const_output_names
)