def check_send_recv(dags):
""" Check result returned by manydags
Check that the from/to jobs exist where the sends say they will
Check that each send/recv has a partner
"""
for m, dag in dags.items():
for outvar, job in dag.items():
if issend(job['fn']) or isrecv(job['fn']):
if issend(job['fn']):
var = job['args'][0]
from_machine = m
to_machine = job['fn'][1]
print "send", to_machine, var, from_machine
if (str(dags[to_machine][var]) !=
str({'fn':('recv', from_machine), 'args':()})):
assert False
return False
if isrecv(job['fn']):
var = outvar
to_machine = m
from_machine = job['fn'][1]
print "recv", to_machine, var, from_machine
if (str({'fn':('send', to_machine), 'args':(var,)}) not in
map(str, dags[from_machine].values())):
assert False
return False
return True
def check_send_recv(dags):
""" Check result returned by manydags
Check that the from/to jobs exist where the sends say they will
Check that each send/recv has a partner
"""
for m, dag in dags.items():
for outvar, job in dag.items():
if issend(job['fn']) or isrecv(job['fn']):
if issend(job['fn']):
var = job['args'][0]
from_machine = m
to_machine = job['fn'][1]
print "send", to_machine, var, from_machine
if (str(dags[to_machine][var]) != str({
'fn': ('recv', from_machine),
'args': ()
})):
assert False
return False
if isrecv(job['fn']):
var = outvar
to_machine = m
from_machine = job['fn'][1]
print "recv", to_machine, var, from_machine
if (str({
'fn': ('send', to_machine),
'args': (var, )
}) not in map(str, dags[from_machine].values())):
assert False
return False
return True
def non_comm_dag(dag):
""" Returns the computational core of dicdag. Cuts off send/recvs.
returns core-of-dicdag, {sent_variables}, {recved_variables}"""
non_comm = {k:v for k,v in dag.items()
if not issend(v['fn']) and not isrecv(v['fn'])}
sent_variables = {v['args'][0] for k, v in dag.items() if issend(v['fn'])}
recvd_variables = {k for k, v in dag.items() if isrecv(v['fn'])}
return non_comm, sent_variables, recvd_variables
def non_comm_dag(dag):
""" Returns the computational core of dicdag. Cuts off send/recvs.
returns core-of-dicdag, {sent_variables}, {recved_variables}"""
non_comm = {
k: v
for k, v in dag.items() if not issend(v['fn']) and not isrecv(v['fn'])
}
sent_variables = {v['args'][0] for k, v in dag.items() if issend(v['fn'])}
recvd_variables = {k for k, v in dag.items() if isrecv(v['fn'])}
return non_comm, sent_variables, recvd_variables
def gpu_dag(dag):
""" The GPU version of a CPU dag - including gpu communication """
nc_dag, sent, recvd = non_comm_dag(dag)
recvs = {
cpu_to_gpu_var(var)[0].clone(): {
'fn': GpuFromHost(),
'args': (var, )
}
for var, it in dag.items() if isrecv(it['fn'])
}
sends = {
it['args'][0]: {
'fn': HostFromGpu(),
'args': (cpu_to_gpu_var(it['args'][0])[0].clone(), )
}
for _, it in dag.items() if issend(it['fn'])
}
def gpu_item((k, v)):
i, op, o = v['args'], v['fn'], (k, )
gi, gop, go = gpu_job(i, op, o)
return (go[0], {'fn': gop, 'args': gi})
gdag = dict(map(gpu_item, nc_dag.items()))
return merge(gdag, recvs, sends)
def ith_output(fn, inputs, idx, old_var):
if issend(fn):
assert len(inputs) == 1 and idx == 0
frommachine = thismachine
_, tomachine = fn
var = inputs[0]
var = theano.tensor.io.send(var,
rankfile[tomachine],
tagfn(frommachine, str(var), tomachine))
var.name = "mpi_token_"+old_var[2:]
return var
if isrecv(fn):
assert len(inputs) == 0 and idx == 0
tomachine = thismachine
_, frommachine = fn
var = theano.tensor.io.recv(known_shape_strings[str(old_var)],
old_var.dtype,
rankfile[frommachine],
tagfn(frommachine, str(old_var),
tomachine))
var.name = old_var.name
return var
return dicdag.theano.theano_dag.ith_output(fn, inputs, idx, old_var)
def ith_output(fn, inputs, idx, old_var):
if issend(fn):
assert len(inputs) == 1 and idx == 0
frommachine = thismachine
_, tomachine = fn
var = inputs[0]
var = theano.tensor.io.send(
var, rankfile[tomachine],
tagfn(frommachine, str(var), tomachine))
var.name = "mpi_token_" + old_var[2:]
return var
if isrecv(fn):
assert len(inputs) == 0 and idx == 0
tomachine = thismachine
_, frommachine = fn
var = theano.tensor.io.recv(
known_shape_strings[str(old_var)], old_var.dtype,
rankfile[frommachine],
tagfn(frommachine, str(old_var), tomachine))
var.name = old_var.name
return var
return dicdag.theano.theano_dag.ith_output(fn, inputs, idx, old_var)
def merge_dags(dags):
""" Merge dags - remove send/recvs between them
input:
dags - dict mapping {machine: dag}
output
Just a single dag
"""
dag = merge(*dags.values())
return {k: v for k,v in dag.items()
if not (issend(v['fn']) and v['fn'][1] in dags)
and not (isrecv(v['fn']) and v['fn'][1] in dags)}
def merge_dags(dags):
""" Merge dags - remove send/recvs between them
input:
dags - dict mapping {machine: dag}
output
Just a single dag
"""
dag = merge(*dags.values())
return {
k: v
for k, v in dag.items() if not (issend(v['fn']) and v['fn'][1] in dags)
and not (isrecv(v['fn']) and v['fn'][1] in dags)
}
def merge_cpu_gpu_dags(cpu_name, cdag, gpu_name, gdag):
""" Merge a cpu and gpu dag - convert the gpu dag first """
if any((issend(v['fn']) or isrecv(v['fn'])) and v['fn'][1] != cpu_name
for v in gdag.values()):
raise Exception("The GPU wants to communicate to someone who isn't the"
" host. We haven't yet built this functionality. TODO")
dag = merge_dags({cpu_name: cdag,
gpu_name: gpu_dag(gdag),
"trans" : gpu_dag_transfers(gdag)})
result = unify_by_name(dag, tuple(variables(merge(cdag,
non_comm_dag(gdag)[0]))))
if any(not isinstance(x, str) and 'gpu' in x.name
for x in inputs_of(result).union(outputs_of(result))):
raise Exception("GPU inputs/outputs")
return result
def merge_cpu_gpu_dags(cpu_name, cdag, gpu_name, gdag):
""" Merge a cpu and gpu dag - convert the gpu dag first """
if any((issend(v['fn']) or isrecv(v['fn'])) and v['fn'][1] != cpu_name
for v in gdag.values()):
raise Exception("The GPU wants to communicate to someone who isn't the"
" host. We haven't yet built this functionality. TODO")
dag = merge_dags({
cpu_name: cdag,
gpu_name: gpu_dag(gdag),
"trans": gpu_dag_transfers(gdag)
})
result = unify_by_name(
dag, tuple(variables(merge(cdag,
non_comm_dag(gdag)[0]))))
if any(not isinstance(x, str) and 'gpu' in x.name
for x in inputs_of(result).union(outputs_of(result))):
raise Exception("GPU inputs/outputs")
return result
def gpu_dag(dag):
""" The GPU version of a CPU dag - including gpu communication """
nc_dag, sent, recvd = non_comm_dag(dag)
recvs = {cpu_to_gpu_var(var)[0].clone(): {'fn': GpuFromHost(),
'args': (var,)}
for var, it in dag.items()
if isrecv(it['fn'])}
sends = {it['args'][0]: {'fn': HostFromGpu(),
'args':(cpu_to_gpu_var(it['args'][0])[0].clone(),)}
for _, it in dag.items()
if issend(it['fn'])}
def gpu_item((k, v)):
i, op, o = v['args'], v['fn'], (k,)
gi, gop, go = gpu_job(i, op, o)
return (go[0], {'fn': gop, 'args': gi})
gdag = dict(map(gpu_item, nc_dag.items()))
return merge(gdag, recvs, sends)
def test_isrecv():
assert isrecv(recv(1,2,3,4))
assert not isrecv(send(1,2,3,4))
def test_isrecv():
assert isrecv(recv(1, 2, 3, 4))
assert not isrecv(send(1, 2, 3, 4))