def test_mpi_send_wait_cmp():
x = theano.tensor.matrix('x')
y = send(x, 1, 11)
z = x + x
waitnode = y.owner
sendnode = y.owner.inputs[0].owner
addnode = z.owner
assert mpi_send_wait_cmp(sendnode, addnode) < 0 # send happens first
assert mpi_send_wait_cmp(waitnode, addnode) > 0 # wait happens last
def test_mpi_schedule():
x = theano.tensor.matrix("x")
y = send(x, 1, 11)
z = x + x
waitnode = y.owner
sendnode = y.owner.inputs[0].owner
addnode = z.owner
f = theano.function([x], [y, z], mode=mpi_mode)
nodes = f.maker.linker.make_all()[-1]
optypes = [MPISend, theano.tensor.Elemwise, MPISendWait]
assert all(isinstance(node.op, optype) for node, optype in zip(nodes, optypes))
def test_send():
x = theano.tensor.matrix('x')
y = send(x, 1, 11)
sendnode = y.owner.inputs[0].owner
assert sendnode.op.dest == 1
assert sendnode.op.tag == 11
stderr.write("mpiexec failed to create a world with two nodes.\n"
"Closing with success message.")
stdout.write("True")
exit(0)
shape = (2, 2)
dtype = 'float32'
scheduler = sort_schedule_fn(*mpi_cmps)
mode = theano.Mode(optimizer=None,
linker=theano.OpWiseCLinker(schedule=scheduler))
if rank == 0:
x = theano.tensor.matrix('x', dtype=dtype)
y = x + 1
send_request = send(y, 1, 11)
z = recv(shape, dtype, 1, 12)
f = theano.function([x], [send_request, z], mode=mode)
xx = np.random.rand(*shape).astype(dtype)
expected = (xx + 1) * 2
_, zz = f(xx)
same = np.linalg.norm(zz - expected) < .001
stdout.write(str(same))
if rank == 1:
"Closing with success message."
)
stdout.write("True")
exit(0)
shape = (2, 2)
dtype = "float32"
scheduler = sort_schedule_fn(*mpi_cmps)
mode = theano.Mode(optimizer=None, linker=theano.OpWiseCLinker(schedule=scheduler))
with config.change_flags(compute_test_value="off"):
if rank == 0:
x = theano.tensor.matrix("x", dtype=dtype)
y = x + 1
send_request = send(y, 1, 11)
z = recv(shape, dtype, 1, 12)
f = theano.function([x], [send_request, z], mode=mode)
xx = np.random.rand(*shape).astype(dtype)
expected = (xx + 1) * 2
_, zz = f(xx)
same = np.linalg.norm(zz - expected) < 0.001
# The parent test will look for "True" in the output
stdout.write(str(same))
if rank == 1:
def test_send():
x = theano.tensor.matrix('x')
y = send(x, 1, 11)
sendnode = y.owner.inputs[0].owner
assert sendnode.op.dest == 1
assert sendnode.op.tag == 11