def _build_chunk_dag(node_str, edge_str):
from mars.tensor.random import TensorRandint
from mars.tensor.arithmetic import TensorTreeAdd
char_dag = DAG()
for s in node_str.split(','):
char_dag.add_node(s.strip())
for s in edge_str.split(','):
l, r = s.split('->')
char_dag.add_edge(l.strip(), r.strip())
chunk_dag = DAG()
str_to_chunk = dict()
for s in char_dag.topological_iter():
if char_dag.count_predecessors(s):
c = TensorTreeAdd(_key=s,
dtype=np.float32()).new_chunk(None,
shape=(10, 10))
inputs = c.op._inputs = [
str_to_chunk[ps] for ps in char_dag.predecessors(s)
]
else:
c = TensorRandint(_key=s,
dtype=np.float32()).new_chunk(None,
shape=(10, 10))
inputs = []
str_to_chunk[s] = c
chunk_dag.add_node(c)
for inp in inputs:
chunk_dag.add_edge(inp, c)
return chunk_dag, str_to_chunk
def testCompose(self):
"""
test compose in build graph and optimize
"""
r"""
graph(@: node, #: composed_node):
@ --> @ --> @ ========> #
"""
chunks = [
TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None)
for n in range(3)
]
graph = DAG()
list(map(graph.add_node, chunks[:3]))
graph.add_edge(chunks[0], chunks[1])
graph.add_edge(chunks[1], chunks[2])
composed_nodes = graph.compose()
self.assertTrue(composed_nodes[0].composed == chunks[:3])
# make the middle one as result chunk, thus the graph cannot be composed
composed_nodes = graph.compose(keys=[chunks[1].key])
self.assertEqual(len(composed_nodes), 0)
r"""
graph(@: node, #: composed_node):
@ @ @ @
\ / \ /
@ --> @ ========> #
/ \ / \
@ @ @ @
"""
chunks = [
TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None)
for n in range(6)
]
graph = DAG()
list(map(graph.add_node, chunks[:6]))
chunks[2].op._inputs = [chunks[0], chunks[1]]
chunks[3].op._inputs = [chunks[2]]
chunks[4].op._inputs = [chunks[3]]
chunks[5].op._inputs = [chunks[3]]
graph.add_edge(chunks[0], chunks[2])
graph.add_edge(chunks[1], chunks[2])
graph.add_edge(chunks[2], chunks[3])
graph.add_edge(chunks[3], chunks[4])
graph.add_edge(chunks[3], chunks[5])
composed_nodes = graph.compose()
self.assertTrue(composed_nodes[0].composed == chunks[2:4])
# to make sure the predecessors and successors of compose are right
# 0 and 1's successors must be composed
self.assertIn(composed_nodes[0], graph.successors(chunks[0]))
self.assertIn(composed_nodes[0], graph.successors(chunks[1]))
# check composed's inputs
self.assertIn(chunks[0].data, composed_nodes[0].inputs)
self.assertIn(chunks[1].data, composed_nodes[0].inputs)
# check composed's predecessors
self.assertIn(chunks[0], graph.predecessors(composed_nodes[0]))
self.assertIn(chunks[1], graph.predecessors(composed_nodes[0]))
# check 4 and 5's inputs
self.assertIn(composed_nodes[0].data,
graph.successors(composed_nodes[0])[0].inputs)
self.assertIn(composed_nodes[0].data,
graph.successors(composed_nodes[0])[0].inputs)
# check 4 and 5's predecessors
self.assertIn(composed_nodes[0], graph.predecessors(chunks[4]))
self.assertIn(composed_nodes[0], graph.predecessors(chunks[5]))
# test optimizer compose
r"""
graph(@: node, S: Slice Chunk, #: composed_node):
@ @ @ @
\ / \ /
@ --> @ --> S ========> # --> S
/ \ / \
@ @ @ @
compose stopped at S, because numexpr don't support Slice op
"""
chunks = [
TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None)
for n in range(6)
]
chunk_slice = TensorSlice().new_chunk([None], None)
graph = DAG()
list(map(graph.add_node, chunks[:6]))
graph.add_node(chunk_slice)
graph.add_edge(chunks[0], chunks[2])
graph.add_edge(chunks[1], chunks[2])
graph.add_edge(chunks[2], chunks[3])
graph.add_edge(chunks[3], chunk_slice)
graph.add_edge(chunk_slice, chunks[4])
graph.add_edge(chunk_slice, chunks[5])
optimizer = NeRuntimeOptimizer(graph)
composed_nodes = optimizer.compose()
self.assertTrue(composed_nodes[0].composed == chunks[2:4])
r"""
graph(@: node, S: Slice Chunk, #: composed_node):
@ --> @ --> S --> @ ========> # --> S --> @
compose stopped at S, because numexpr don't support Slice op
"""
chunks = [
TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None)
for n in range(4)
]
graph = DAG()
list(map(graph.add_node, chunks[:3]))
graph.add_node(chunk_slice)
graph.add_edge(chunks[0], chunks[1])
graph.add_edge(chunks[1], chunk_slice)
graph.add_edge(chunk_slice, chunks[2])
optimizer = NeRuntimeOptimizer(graph)
composed_nodes = optimizer.compose()
self.assertTrue(composed_nodes[0].composed == chunks[:2])
self.assertTrue(len(composed_nodes) == 1)
r"""
graph(@: node, S: Slice Chunk, #: composed_node):
@ --> @ --> S --> @ --> @ ========> # --> S --> #
compose stopped at S, because numexpr don't support Slice op
"""
chunks = [
TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None)
for n in range(4)
]
graph = DAG()
list(map(graph.add_node, chunks[:4]))
graph.add_node(chunk_slice)
graph.add_edge(chunks[0], chunks[1])
graph.add_edge(chunks[1], chunk_slice)
graph.add_edge(chunk_slice, chunks[2])
graph.add_edge(chunks[2], chunks[3])
optimizer = NeRuntimeOptimizer(graph)
composed_nodes = optimizer.compose()
self.assertTrue(composed_nodes[0].composed == chunks[:2])
self.assertTrue(composed_nodes[1].composed == chunks[2:4])
def testDAG(self):
r"""
1 --- 4
2 --- 6
\ /
5
/
3
"""
dag = DAG()
[dag.add_node(i) for i in range(1, 7)]
dag.add_edge(1, 4)
dag.add_edge(2, 6)
dag.add_edge(2, 5)
dag.add_edge(5, 6)
dag.add_edge(3, 5)
with self.assertRaises(KeyError):
dag.add_edge(1, 10)
with self.assertRaises(KeyError):
dag.add_edge(10, 1)
self.assertEqual(set(dag[2]), {5, 6})
self.assertEqual(list(dag.topological_iter()), [3, 2, 5, 6, 1, 4])
self.assertEqual(list(dag.dfs()), [3, 2, 5, 6, 1, 4])
self.assertEqual(list(dag.bfs()), [1, 2, 3, 4, 5, 6])
dag.add_edge(6, 1)
dag.add_edge(1, 2)
with self.assertRaises(KeyError):
for _ in dag.iter_predecessors(-1):
pass
with self.assertRaises(KeyError):
for _ in dag.iter_successors(-1):
pass
self.assertRaises(GraphContainsCycleError,
lambda: list(dag.topological_iter()))
dag.remove_edge(2, 5)
self.assertFalse(dag.has_successor(2, 5))
with self.assertRaises(KeyError):
dag.remove_edge(2, 5)
rev_dag = dag.build_reversed()
for n in dag:
self.assertIn(n, rev_dag)
self.assertTrue(
all(
rev_dag.has_successor(n, pred)
for pred in dag.predecessors(n)))
undigraph = dag.build_undirected()
for n in dag:
self.assertIn(n, undigraph)
self.assertTrue(
all(
undigraph.has_predecessor(pred, n)
for pred in dag.predecessors(n)))
self.assertTrue(
all(
undigraph.has_successor(n, pred)
for pred in dag.predecessors(n)))
dag_copy = dag.copy()
for n in dag:
self.assertIn(n, dag_copy)
self.assertTrue(
all(
dag_copy.has_successor(pred, n)
for pred in dag_copy.predecessors(n)))
