import torch.fx as fx
import copy
import torch
import math
from typing import Callable, List
from functools import wraps, partial
from dataclasses import dataclass
from .compile_utils import get_placeholders, get_outputs
class ConcreteProp(torch.fx.Interpreter):
def run_node(self, n):
result = super().run_node(n)
found_tensor = False
def extract_tensor_meta(obj):
if isinstance(obj, torch.Tensor):
nonlocal found_tensor
found_tensor = True
return obj
else:
return obj
from torch.fx.node import map_aggregate
concrete_value = map_aggregate(result, extract_tensor_meta)
if found_tensor:
n.meta['concrete_value'] = concrete_value
return result
def propagate(self, *args):
return super().run(*args)
# inplace modifies node/inps
def _convert_node_to_placeholder(node, inps):
if node.op == 'output' or node.op == "placeholder":
return
node.op = 'placeholder'
node.args = ()
node.kwargs = {}
node.target = node.name
concrete_val = node.meta.get('concrete_value', None)
if isinstance(concrete_val, torch.Tensor):
inps.append(concrete_val)
else:
inps.append(torch.zeros(()))
for tuple_user in list(node.users):
_convert_node_to_placeholder(tuple_user, inps)
def dump_state(fx_g, inps):
print(f"""
# Working Repro with {len(fx_g.graph.nodes)} nodes
inps = {[(i.shape, i.dtype, i.device.type) for i in inps]}
inps = [torch.zeros(())] + [torch.ones(shape, dtype=dtype, device=device) for (shape, dtype, device) in inps]
{fx_g.code}
""")
@dataclass
class ReproState:
graph: fx.Graph
inps: List[torch.Tensor]
def minifier(fail_f: fx.GraphModule, inps, module_fails, dump_state: Callable = dump_state):
"""
Minimizes a FX graph with given inputs, such that the resulting FX graph still returns True for module_fails.
Does 2 main strategies:
1. Truncates suffix: Removes some suffix from the graph and sets a new output.
2. Delta Debugging: Tries replacing half of the graph with inputs. If fails,
tries replacing quarter of the graph, etc.
>>> # xdoctest: +SKIP(failing)
>>> failing_function = fx.symbolic_trace(f)
>>> minimize(failing_function, [torch.randn(5)], lambda fx_g, inps: fx_g(*inps))
note: module_fails returns True if it fails.
"""
failing_graph = fail_f.graph
cur_size = len(failing_graph.nodes)
num_queries = 0
def deepcopy_fx_graph(fx_graph):
return fx.GraphModule(fail_f, copy.deepcopy(fx_graph)).graph
def graph_fails(graph, inps):
nonlocal num_queries
graph = copy.deepcopy(graph)
num_queries += 1
mod = fx.GraphModule(fail_f, graph)
mod.graph.lint()
return module_fails(mod, inps)
ConcreteProp(fail_f).propagate(*inps)
if not graph_fails(failing_graph, inps):
raise RuntimeError("Input graph did not fail the tester")
print(f"Started off with {cur_size} nodes")
def _register_strategy(strategy: Callable, name: str):
@wraps(strategy)
def new_func(old_state: ReproState, granularity=1):
print()
print(f"Strategy: {name} (G: {granularity}) ({len(old_state.graph.nodes)} nodes, {len(old_state.inps)} inputs)")
new_state = strategy(deepcopy_fx_graph(old_state.graph), list(old_state.inps), granularity)
if new_state is not None:
new_nodes = len(new_state.graph.nodes)
old_nodes = len(old_state.graph.nodes)
new_inps = len(new_state.inps)
old_inps = len(old_state.inps)
new_outs = len(get_outputs(new_state.graph))
old_outs = len(get_outputs(old_state.graph))
progress_made = False
if new_nodes < old_nodes:
progress_made = True
print(f"SUCCESS: Went from {old_nodes} to {new_nodes} nodes")
if new_inps > old_inps:
progress_made = True
print(f"SUCCESS: Went from {old_inps} to {new_inps} inputs")
if new_outs < old_outs:
progress_made = True
print(f"SUCCESS: Went from {old_outs} to {new_outs} outputs")
if not progress_made:
raise RuntimeError("Success raised but no progress made?")
if not graph_fails(new_state.graph, new_state.inps):
print("WARNING: Something went wrong, not applying this minification")
return None
return new_state
else:
print(f"FAIL: {name}")
return None
return new_func
def register_strategy(name: str):
return partial(_register_strategy, name=name)
@register_strategy("Truncate suffix")
def remove_suffix(cur_graph, cur_inps, granularity):
tested = set()
new_graph = fx.Graph()
env = {}
for idx, node in enumerate(cur_graph.nodes):
new_node = new_graph.node_copy(node, lambda x: env[x])
if node.op not in ['placeholder', 'output']:
# If idx is divisible by (granularity * 2), it would have been checked already.
if idx % granularity == 0 and (idx % (granularity * 2) != 0) and idx not in tested:
output_node = new_graph.output((new_node,))
if len(new_graph.nodes) < len(cur_graph.nodes) and graph_fails(new_graph, cur_inps):
return ReproState(new_graph, cur_inps)
else:
tested.add(idx)
new_graph.erase_node(output_node)
env[node] = new_node
return None
@register_strategy("Remove outputs")
def remove_outputs(cur_graph, cur_inps, granularity):
granularity = max(1, granularity // 2)
for idx, node in enumerate(cur_graph.nodes):
node.idx = idx
if node.op == 'output':
output = node
break
output_args = sorted(output.args[0], key=lambda x: x.idx if isinstance(x, fx.Node) else int(1e9))
if len(output_args) == 1:
return None
for idx in range(0, len(output_args), granularity):
output.args = (output_args[:idx] + output_args[idx + granularity:],)
if graph_fails(cur_graph, cur_inps):
return ReproState(cur_graph, cur_inps)
return None
def remove_unused_inputs_unchecked(cur_state: ReproState):
cur_graph = cur_state.graph
cur_inps = cur_state.inps
ph_nodes = get_placeholders(cur_graph)
assert len(ph_nodes) == len(cur_inps)
new_inps = []
for idx in range(len(ph_nodes)):
if len(ph_nodes[idx].users) == 0:
cur_graph.erase_node(ph_nodes[idx])
else:
new_inps.append(cur_inps[idx])
if len(new_inps) < len(cur_inps):
return ReproState(cur_graph, new_inps)
return None
def remove_unused_inputs_checked(cur_state: ReproState):
new_state = remove_unused_inputs_unchecked(cur_state)
if new_state is not None and graph_fails(new_state.graph, new_state.inps):
return new_state
return None
def _remove_unused_wrapper(cur_graph, cur_inps, granularity):
return remove_unused_inputs_checked(ReproState(cur_graph, cur_inps))
remove_unused_inputs = register_strategy("Remove unused inputs")(_remove_unused_wrapper)
@register_strategy("Eliminate dead code")
def eliminate_dead_code(cur_graph, cur_inps, granularity):
if cur_graph.eliminate_dead_code() and graph_fails(cur_graph, cur_inps):
return ReproState(cur_graph, cur_inps)
return None
def _consolidate_placeholders(cur_graph):
new_graph = fx.Graph()
env = {}
for node in cur_graph.nodes:
if node.op == 'placeholder':
new_node = new_graph.node_copy(node, lambda x: env[x])
env[node] = new_node
for node in cur_graph.nodes:
if node.op != 'placeholder':
new_node = new_graph.node_copy(node, lambda x: env[x])
env[node] = new_node
return new_graph
@register_strategy("Delta Debugging")
def delta_debugging(cur_graph: fx.Graph, cur_inps, granularity):
num_nodes = len(cur_graph.nodes)
for start_range in range(0, num_nodes, granularity):
is_removing = False
new_graph = deepcopy_fx_graph(cur_graph)
new_inps = cur_inps[:]
end_range = min(num_nodes, start_range + granularity)
for idx in range(start_range, end_range):
new_node = list(new_graph.nodes)[idx]
if new_node.op not in ['placeholder', 'output']:
is_removing = True
_convert_node_to_placeholder(new_node, new_inps)
if not is_removing:
continue
new_graph = _consolidate_placeholders(new_graph)
new_state = remove_unused_inputs_unchecked(ReproState(new_graph, new_inps))
if new_state is None:
new_state = ReproState(new_graph, new_inps)
if graph_fails(new_state.graph, new_state.inps):
return ReproState(new_state.graph, new_state.inps)
return None
failing_state = ReproState(failing_graph, inps)
def try_granularity(failing_state, granularity, use_non_granular):
print(f"Trying granularity {granularity}")
strategies = []
num_nodes = len(failing_state.graph.nodes)
num_outputs = len(get_outputs(failing_state.graph))
if num_outputs > num_nodes // 2:
strategies += [remove_outputs]
if use_non_granular:
strategies += [eliminate_dead_code, remove_unused_inputs]
strategies += [remove_suffix, delta_debugging]
for strategy in strategies:
new_state = strategy(failing_state, granularity)
if new_state is not None:
return new_state
return None
while True:
dump_state(fx.GraphModule(fail_f, failing_state.graph), failing_state.inps)
granularity = int(2**(math.floor(math.log2(len(failing_state.graph.nodes)))))
new_state = try_granularity(failing_state, granularity, use_non_granular=True)
if new_state is not None:
failing_state = new_state
continue
granularity //= 2
has_progress = False
while granularity >= 1:
new_state = try_granularity(failing_state, granularity, use_non_granular=False)
if new_state is not None:
failing_state = new_state
has_progress = True
break
granularity //= 2
if has_progress:
continue
new_state = remove_outputs(failing_state, 1)
if new_state is not None:
failing_state = new_state
continue
break
if not graph_fails(failing_state.graph, failing_state.inps):
raise RuntimeError("Uh oh, something went wrong :( Final graph is not failing")
print(f"Made {num_queries} queries")
failing_fx = fx.GraphModule(fail_f, failing_state.graph)
dump_state(failing_fx, failing_state.inps)
print("Wrote minimal repro out to repro.py")
return failing_fx, failing_state.inps