nengo / nengo-loihi

Run Nengo models on Intel's Loihi chip
https://www.nengo.ai/nengo-loihi/
Other
35 stars 12 forks source link

Split blocks automatically #264

Closed hunse closed 4 years ago

hunse commented 4 years ago

Loihi is organized into cores with a fixed number of compartments on each, and since the start we've required users to manually break their model into Ensembles that will each fit on one core.

This PR automates splitting larger ensembles to fit across multiple cores. This allows users to create the model in terms of structures that work well conceptually, and worry less about how that is going to map to Loihi.

There are two ways of using this functionality. One is to let nengo-loihi figure things out itself, in which case it simply splits large ensembles sequentially (putting the first N neurons on one core, the next N on the next core, etc.). This works well for NEF or other fully-connected Ensembles that have a fairly uniform structure in terms of input and output connections.

The second way to use block splitting is to provide instructions on how you (as a user) want an ensemble to be split. This is done via the full_shape and block_shape config options.

with nengo.Network() as net:
    a = nengo.Ensemble(120)
    net.config[a].full_shape = (6, 5, 4)
    net.config[a].block_shape = (3, 2, 3)

The block_shape specifies the shape that a single block (i.e. one core) will represent. The maximum number of compartments on that core is the product of all numbers of the shape. We then tile that shape to fill the full shape. So in the above example, we'll have 2 cores in the first dimension (since 6 \ 3 = 2, where \ represents ceiling division ceil(a / b)), 3 cores in the second dimension (5 \ 2 = 3), and 2 cores in the third (4 \ 3 = 2). The total number of cores is 2 * 3 * 2 = 12, and the layout of the cores is (2, 3, 2). We then "rebalance" the block_shape so that it is as uniform as possible across cores, given this layout, by taking the ceiling division of each element of the full shape by the corresponding number of cores in that dimension: (6, 5, 4) \ (2, 3, 2) = (3, 2, 2). You can see this is close to the original block shape, but with the last dimension being 2 instead of 3. What's happened is that in the first dimension, we'll have 2 cores of length 3 go evenly into 6, so everything is already balanced there. In the second dimension, we'll have 2 cores of length 2 and one of length 1 to make up 5. This isn't perfectly balanced, but there's no way to make one core shorter and another longer to have it be better. But in the last dimension, with the original block shape we would have one core of length 3 and one of length 1 to make up the 4. We can instead have two cores of length 2, which is more balanced (and doesn't use any more cores), so we do that instead.

This PR also adds a number of features that we needed for a recent project.

Based on #261.

TODO:

tbekolay commented 4 years ago

Note to self: see if the spiking MNIST example can be simplified using this PR.