The computation of the Kalman filter over two computers using MPI-enabled Theano.

kalman_mpi.py

import theano
from theano.tensor.io import send, recv, mpi_cmps
import theano.sandbox.linalg as linalg
from theano.gof.sched import sort_schedule_fn
from time import time

dot = theano.tensor.dot
dtype = 'float32'
n = 500
run = False

# Set up a linker that orders nodes to overlap computation and communication
mpi_scheduler = sort_schedule_fn(*mpi_cmps)
mpi_linker = theano.OpWiseCLinker(schedule=mpi_scheduler)
mpi_mode = theano.Mode(linker=mpi_linker)

# initialize MPI
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
rank = comm.Get_rank()

if rank == 0 or not run:
    # Create some input variables
    mu      = theano.tensor.matrix('mu')
    Sigma   = theano.tensor.matrix('Sigma')
    H       = theano.tensor.matrix('H')
    R       = theano.tensor.matrix('R')
    data    = theano.tensor.matrix('data')
    input_shapes = { mu:     (n, 1),
                     Sigma:  (n, n),
                     H:      (n, n),
                     R:      (n, n),
                     data:   (n, 1)}

    # Some intermediate variables
    A = dot(Sigma, H.T)
    B = R + dot(H, dot(Sigma, H.T))

    new_mu    = mu + dot(A, linalg.solve(B, dot(H, mu) - data))
    new_mu.name = "updated_mu"

    # Send data to 1
    receipts = send(H, 1, 1), send(B, 1, 2), send(Sigma, 1, 3), send(A, 1, 4)

    # Get back the work that 1 did
    new_Sigma = recv((n, n), dtype, 1, 5)

    # Compile
    inputs_0 = (mu, Sigma, H, R, data)
    outputs_0 = (new_mu, new_Sigma) + receipts
    f0 = theano.function(inputs_0, outputs_0, mode=mpi_mode)
    nodes0 = f0.maker.linker.make_all()[-1] # for debug


    if run:
        # Generate random inputs
        numeric_inputs = [np.random.rand(*input_shapes[inp]).astype(dtype)
                                                for inp in inputs_0]

        a, b, _, _, _, _ = f0(*numeric_inputs) # warm start
        # Run and time
        comm.barrier()
        starttime = time()
        a, b, _, _, _, _ = f0(*numeric_inputs)
        comm.barrier()
        endtime = time()

        print endtime - starttime

if rank == 1 or not run:
    # Receive some data from 0
    H       = recv((n, n), dtype, 0, 1)
    B       = recv((n, n), dtype, 0, 2)
    Sigma   = recv((n, n), dtype, 0, 3)
    A       = recv((n, n), dtype, 0, 4)

    # Do some computation
    new_Sigma = Sigma - dot(dot(A, linalg.solve(B, H)), Sigma)
    new_Sigma.name = "updated_Sigma"

    # Send it back to 0
    receipt = send(new_Sigma, 0, 5)

    # compile locally using Theano
    inputs_1 = ()
    outputs_1 = (receipt, )
    f1 = theano.function(inputs_1, outputs_1, mode=mpi_mode)
    nodes1 = f1.maker.linker.make_all()[-1] # for debug

    if run:
        _ = f1() # warm start
        comm.barrier()
        _ = f1()
        comm.barrier()

This is the resultant ordering of the nodes. Notice how Theano orders sends, receives, and waits in order to overlap communication and computation and block as little as possible.

>>> nodes0
[MPISend{dest: 1, tag: 1}(H),
 MPISend{dest: 1, tag: 3}(Sigma),
 MPIRecv{source: 1, tag: 5, shape: (500, 500), dtype: float32}(),
 InplaceDimShuffle{1,0}(H),
 _dot22(Sigma, H.T),
 MPISend{dest: 1, tag: 4}(_dot22.0),
 Gemm{no_inplace}(R, TensorConstant{1.0}, H, _dot22.0, TensorConstant{1.0}),
 MPISend{dest: 1, tag: 2}(Gemm{no_inplace}.0),
 Gemm{no_inplace}(data, TensorConstant{1.0}, H, mu, TensorConstant{-1.0}),
 Solve{('general', False, False, False)}(Gemm{no_inplace}.0, Gemm{no_inplace}.0),
 Gemm{no_inplace}(mu, TensorConstant{1.0}, _dot22.0, Solve{('general', False, False, False)}.0, TensorConstant{1.0}),
 MPISendWait(MPISend{dest: 1, tag: 1}.0),
 MPISendWait(MPISend{dest: 1, tag: 2}.0),
 MPISendWait(MPISend{dest: 1, tag: 3}.0),
 MPISendWait(MPISend{dest: 1, tag: 4}.0),
 MPIRecvWait(MPIRecv{source: 1, tag: 5, shape: (500, 500), dtype: float32}.0, MPIRecv{source: 1, tag: 5, shape: (500, 500), dtype: float32}.1)]

>>> nodes1
[MPIRecv{source: 0, tag: 1, shape: (500, 500), dtype: float32}(),
 MPIRecv{source: 0, tag: 2, shape: (500, 500), dtype: float32}(),
 MPIRecv{source: 0, tag: 3, shape: (500, 500), dtype: float32}(),
 MPIRecv{source: 0, tag: 4, shape: (500, 500), dtype: float32}(),
 MPIRecvWait(MPIRecv{source: 0, tag: 1, shape: (500, 500), dtype: float32}.0, MPIRecv{source: 0, tag: 1, shape: (500, 500), dtype: float32}.1),
 MPIRecvWait(MPIRecv{source: 0, tag: 2, shape: (500, 500), dtype: float32}.0, MPIRecv{source: 0, tag: 2, shape: (500, 500), dtype: float32}.1),
 Solve{('general', False, False, False)}(MPIRecvWait.0, MPIRecvWait.0),
 MPIRecvWait(MPIRecv{source: 0, tag: 4, shape: (500, 500), dtype: float32}.0, MPIRecv{source: 0, tag: 4, shape: (500, 500), dtype: float32}.1),
 _dot22(MPIRecvWait.0, Solve{('general', False, False, False)}.0),
 MPIRecvWait(MPIRecv{source: 0, tag: 3, shape: (500, 500), dtype: float32}.0, MPIRecv{source: 0, tag: 3, shape: (500, 500), dtype: float32}.1),
 _dot22(_dot22.0, MPIRecvWait.0),
 Elemwise{Sub{output_types_preference=transfer_type{0}}}[(0, 0)](MPIRecvWait.0, _dot22.0),
 MPISend{dest: 0, tag: 5}(updated_Sigma),
 MPISendWait(MPISend{dest: 0, tag: 5}.0)]

MPISend and MPIRecv are asynchronous. They set everything up but take no time themselves. They are done as soon as possible.
MPI*Waits block until the transfer is complete. They are done as late as possible.
MPI transfers are given unique identifiers (tags). Tags are used to break ties. This helps to prevent deadlocks.