def test_compute_topological_order():
from pytools.graph import compute_topological_order, CycleError
empty = {}
assert compute_topological_order(empty) == []
disconnected = {1: [], 2: [], 3: []}
assert len(compute_topological_order(disconnected)) == 3
line = list(zip(range(10), ([i] for i in range(1, 11))))
import random
random.seed(0)
random.shuffle(line)
expected = list(range(11))
assert compute_topological_order(dict(line)) == expected
claw = {1: [2, 3], 0: [1]}
assert compute_topological_order(claw)[:2] == [0, 1]
repeated_edges = {1: [2, 2], 2: [0]}
assert compute_topological_order(repeated_edges) == [1, 2, 0]
self_cycle = {1: [1]}
with pytest.raises(CycleError):
compute_topological_order(self_cycle)
cycle = {0: [2], 1: [2], 2: [3], 3: [4, 1]}
with pytest.raises(CycleError):
compute_topological_order(cycle)
def test_prioritzed_topological_sort_examples():
from pytools.graph import compute_topological_order
keys = {"a": 4, "b": 3, "c": 2, "e": 1, "d": 4}
dag = {"a": ["b", "c"], "b": [], "c": ["d", "e"], "d": [], "e": []}
assert compute_topological_order(
dag, key=keys.get) == ["a", "c", "e", "b", "d"]
keys = {"a": 7, "b": 2, "c": 1, "d": 0}
dag = {
"d": set("c"),
"b": set("a"),
"a": set(),
"c": set("a"),
}
assert compute_topological_order(dag, key=keys.get) == ["d", "c", "b", "a"]
def make_partition(self, outputs: DictOfNamedArrays) -> GraphPartition:
rewritten_outputs = {
name: self(expr) for name, expr in outputs._data.items()}
pid_to_output_names: Dict[PartId, Set[str]] = {
pid: set() for pid in self.seen_part_ids}
pid_to_input_names: Dict[PartId, Set[str]] = {
pid: set() for pid in self.seen_part_ids}
var_name_to_result = self.var_name_to_result.copy()
for out_name, rewritten_output in rewritten_outputs.items():
out_part_id = self._get_part_id(outputs._data[out_name])
pid_to_output_names.setdefault(out_part_id, set()).add(out_name)
var_name_to_result[out_name] = rewritten_output
# Mapping of nodes to their successors; used to compute the topological order
pid_to_needing_pids: Dict[PartId, Set[PartId]] = {
pid: set() for pid in self.seen_part_ids}
pid_to_needed_pids: Dict[PartId, Set[PartId]] = {
pid: set() for pid in self.seen_part_ids}
for (pid_target, pid_dependency), var_names in \
self.part_pair_to_edges.items():
pid_to_needing_pids[pid_dependency].add(pid_target)
pid_to_needed_pids[pid_target].add(pid_dependency)
for var_name in var_names:
pid_to_output_names[pid_dependency].add(var_name)
pid_to_input_names[pid_target].add(var_name)
from pytools.graph import compute_topological_order, CycleError
try:
toposorted_part_ids = compute_topological_order(pid_to_needing_pids)
except CycleError:
raise PartitionInducedCycleError
return GraphPartition(
parts={
pid: GraphPart(
pid=pid,
needed_pids=frozenset(pid_to_needed_pids[pid]),
user_input_names=frozenset(
self.pid_to_user_input_names.get(pid, set())),
partition_input_names=frozenset(pid_to_input_names[pid]),
output_names=frozenset(pid_to_output_names[pid]),
)
for pid in self.seen_part_ids},
var_name_to_result=var_name_to_result,
toposorted_part_ids=toposorted_part_ids)
def _toposort_of_subset_of_insns(kernel, subset_insns):
"""
Returns a :class:`list` of insn ids which is a topological sort of insn
deps in *subset_insns*.
:arg subset_insns: a :class:`frozenset` of insn ids that are a subset of
kernel over which we wish to compute the topological sort.
"""
dag = {
insn_id: set(kernel.id_to_insn[insn_id].depends_on
& subset_insns)
for insn_id in subset_insns
}
from pytools.graph import compute_topological_order
return compute_topological_order(dag)[::-1]
def test_prioritzed_topological_sort():
import random
from pytools.graph import compute_topological_order
rng = random.Random(0)
def generate_random_graph(nnodes):
graph = {i: set() for i in range(nnodes)}
for i in range(nnodes):
# to avoid cycles only consider edges node_i->node_j where j > i.
for j in range(i + 1, nnodes):
# Edge probability 4/n: Generates decently interesting inputs.
if rng.randint(0, nnodes - 1) <= 2:
graph[i].add(j)
return graph
nnodes = rng.randint(40, 100)
rev_dep_graph = generate_random_graph(nnodes)
dep_graph = {i: set() for i in range(nnodes)}
for i in range(nnodes):
for rev_dep in rev_dep_graph[i]:
dep_graph[rev_dep].add(i)
keys = [rng.random() for _ in range(nnodes)]
topo_order = compute_topological_order(rev_dep_graph, key=keys.__getitem__)
for scheduled_node in topo_order:
nodes_with_no_deps = {
node
for node, deps in dep_graph.items() if len(deps) == 0
}
# check whether the order is a valid topological order
assert scheduled_node in nodes_with_no_deps
# check whether priorites are upheld
assert keys[scheduled_node] == min(keys[node]
for node in nodes_with_no_deps)
# 'scheduled_node' is scheduled => no longer a dependency
dep_graph.pop(scheduled_node)
for node, deps in dep_graph.items():
deps.discard(scheduled_node)
assert len(dep_graph) == 0
def preprocess(outputs: DictOfNamedArrays) -> PreprocessResult:
"""Preprocess a computation for code generation."""
from pytato.transform import copy_dict_of_named_arrays, get_dependencies
# {{{ compute the order in which the outputs must be computed
# semantically order does not matter, but doing a toposort ordering of the
# outputs leads to a FLOP optimal choice
from pytools.graph import compute_topological_order
deps = get_dependencies(outputs)
# only look for dependencies between the outputs
deps = {
name: (val & frozenset(outputs.values()))
for name, val in deps.items()
}
# represent deps in terms of output names
output_to_name = {output: name for name, output in outputs.items()}
dag = {
name: (frozenset([output_to_name[output]
for output in val]) - frozenset([name]))
for name, val in deps.items()
}
output_order: List[str] = compute_topological_order(dag)[::-1]
# }}}
mapper = CodeGenPreprocessor(Namespace())
new_outputs = copy_dict_of_named_arrays(outputs, mapper)
return PreprocessResult(outputs=new_outputs,
compute_order=tuple(output_order),
bound_arguments=mapper.bound_arguments)