[Nengo-Loihi]: Adding NxSDK code to Nengo-Loihi/Nengo

Hello everyone,

In my understanding, Nengo-Loihi is sort of a wrapper (and perhaps more) over NxSDK (which provides low level python APIs for programming on Intel’s Loihi chips). I was wondering if there are tutorials for writing custom NxSDK code and wrapping it in Nengo-Loihi (or even better, Nengo if it supports). On the INRC forum, I believe there are some tutorials for using NxSDK available, but what about using NxSDK in conjunction with Nengo-Loihi/Nengo? Please let me know.

Thanks!

Unfortunately, we don’t really have any tutorials for this currently. So you’ll probably have to do some poking around in the NengoLoihi code to figure out exactly how we’re configuring things with NxSDK, and then figure out how what you want to add can work with this.

Even more unfortunately, all of our NxSDK calls are obfuscated, because there were concerns with having the NxSDK API exposed in a searchable way online. This makes it much harder to understand at a glance exactly how we’re configuring NxSDK.

Here’s a brief skeleton of how you might set things up:

# create outside `with` block, because we don't want to connect to the board yet
sim = nengo_loihi.Simulator(net)

# get our `nengo_loihi.hardware.interface.HardwareInterface` object
hw_interface = sim.sims["loihi"]

# get the NxSDK board object
nxsdk_board = hw_interface.nxsdk_board

# get our `nengo_loihi.hardware.nxsdk_objects.Board` object
board = hw_interface.board

# do configuration here

with sim:
    # run simulation here with e.g. sim.run

Here, our board object is designed to mirror the nxsdk_board object, in that the index of a chip/core in the board object matches the index it will have in the nxsdk_board object. You can use board.find_block and board.find_synapse to figure out where a nengo_loihi.block.LoihiBlock or nengo_loihi.block.Synapse are located in terms of chip/core/compartments. You can then use this information when configuring nxsdk_board. You can use the sim.model.objs dictionary to map from Nengo objects (e.g. Ensemble) to NengoLoihi objects (e.g. LoihiBlock).

Hello @Eric, thanks for your response. Looks like it’s going to be a high learning curve in the absence of tutorials; and moreover it appears to me that one needs to familiarize him/herself with NengoLoihi in general… before tinkering with NxSDK APIs using the obfuscation interface of NengoLoihi. For now, I think I should rest my case here about writing custom NxSDK code and integrating it with NengoLoihi, unless my plans change.

Just curious, in case I decide to… how should I progress? Let’s suppose I want to develop a custom network of few connected individual neurons which will execute low level NxSDK APIs (e.g. its join functions), do I straight-ahead write the network in NxSDK and then find out how to wrap it in NengoLoihi, or do I use the obfuscation interface of NengoLoihi to create such a custom network using NxSDK and make those low level calls?

I guess in your above example, you do talk about do configuration here which to me seems that we have to use the board object to make low level NxSDK calls.

It’s really hard to suggest the best way to proceed without a concrete goal or example of exactly what you want to accomplish.

Part of it depends on how close your goal is to what can already be done with NengoLoihi. If it is fairly similar, then you could make your model in NengoLoihi, and then go in and just tweak the parts you need on the NxSDK board object. For example, you could build your network as normal and then go in and add delays to particular synapses (something that we don’t currently support in NengoLoihi).

One thing we’ve done before is writing custom SNIPs for a particular project. This is not too hard to do with NengoLoihi; the main thing you have to do is figure out where the neurons (or whatever else) you want to interact with are located on the board, and then you can target those things in your SNIP, and use the NxSDK interface to add the SNIP.

On the other hand, if what you want to do is very far from what can be done in NengoLoihi, then it might be better to write it separately in NxSDK. The difficulty with this is that NxSDK requires you to configure the number of chips/cores/synapses when you create the NxSDK board object; we do this here. So you’d somehow have to hack that to leave extra space for whatever you want to build, and then build into that space. One way to hack it might be to have one or more Ensembles that do nothing and aren’t connected to anything, which you can then find and overwrite on the board with your custom stuff.

Hello @Eric, I see that it is getting more complex. I don’t have any experience of writing custom SNIPs. BTW, to explain what I want to achieve at the end of the day, I would like to expand on the following:

I want to leverage the “join operations” that have been mentioned in the u_join_op_in_multi_compartment_neuron.ipynb tutorial file in the INRC VM nxsdk directory. A short description about the “join operation” has also been mentioned in the Loihi paper as follows:

The compartments’ state variables are combined in a configurable 
manner by programming different join functions for each 
compartment junction.

I am not sure if you are aware of them, therefore mentioning these following details (sorry for the blurt if you already are). The compartmental neurons in a neurocore are arranged in binary tree structure on the Loihi chip, and the parent neuron “joins” the output from its two individual child neurons. Different “join” ops could be “ADD”, “MAX”, “MIN”, etc. Now, the ipynb file I mentioned above, has shown how to do different “join” ops using the NxSDK APIs; I was wondering to do the same, albeit not limiting to just one instance of a join op, rather create a layer of many such instances. The input to and output from the “join” op layer will be from a previous Conv and to the next Conv layer respectively; thus I will have to integrate my custom NxSDK code with NengoLoihi.

In light of above, what could be the best way to progress? Is it possible to create such a layer of “join” ops directly with NengoLoihi or do I write it in NxSDK API explicitly and then somehow try to integrate it with the rest of my network in NengoLoihi?

My instinct would be to create an ensemble, connect in to a set of input neurons in that ensemble (which will be represented on Loihi as single compartments), connect out from a separate set of output neurons (which will create output compartments), then go in and manually connect those input/output compartments.

Here’s a sketch:

with nengo.Network() as net:
    n = 10  # number of input/output compartments
    phases = np.linspace(0, 2*np.pi, n)
    inp = nengo.Node(lambda t: np.sin(t + phases))

    ens = nengo.Ensemble(2 * n, 1)
    out = nengo.Node(size_in=n)
    probe = nengo.Probe(out)

    nengo.Connection(inp, ens.neurons[:n])
    nengo.Connection(ens.neurons[n:], out)

sim = nengo_loihi.Simulator(net)
hw_interface = sim.sims["loihi"]
nxsdk_board = hw_interface.nxsdk_board
board = hw_interface.board

blocks = sim.model.objs[ens]
assert len(blocks) == 1
block = blocks[0]

# get the nengo_loihi.block.Synapse
synapse = sim.model.objs[input_conn]["weights"]

chip_idx, core_idx, block_idx, compartment_idxs, _ = board.find_block(block)
nxsdk_core = nxsdk_board.n2Chips[chip_idx].n2CoresAsList[core_idx]

synapse_unique_idxs = np.unique(synapse.indices)
input_compartment_idxs = [compartment_idxs[i] for i in synapse_unique_idxs]
output_compartment_idxs = [i for i in compartment_idxs if i not in input_compartment_idxs]

# do configuration here using nxsdk_core, and input/output compartment idxs
# NengoLoihi currently only puts one block per core, so you should be able to use
# any compartments not in `compartment_idxs` as extra compartments that you can
# configure yourself. 

Hello @Eric, thanks for the explanation and the sample code. About the nxsdk_board and board object, I guess, they will be created only when a physical board would be present. Right?

Also about configuring/manually connecting the input/output neurons, I would be required to makes sure that the input neurons are twice in numbers compared to the number of output neurons and moreover, each group of two input neurons should be the children nodes of the parent output neuron. Will it be possible with your method of obtaining the input_compartment_idxs and output_compartment_idxs? I am sorry but I wasn’t able to follow the code quite clearly after the line # get the nengo_loihi.block.Synapse.

So essentially each nengo_loihi.block.Synapse is responsible for one nengo.Connection, and on the Loihi board connects input axons to compartments using synaptic weights. If you want your input compartments to each be coming from different connections, then you can just repeat what I have above to get separate input_compartment_idxs for each connection. If, on the other hand, all your input compartments are being fed by the same connection, then you’ll have to take the single list of input_compartment_idxs and group them yourself (I assume in something like a pair of input compartments for each output compartment).

As for your first question, yes, this is really only possible when you have a board to connect to. Technically, all that’s required to create the board and nxsdk_board is for NxSDK to be installed, so it would be possible to run that on your machine without a board if you have NxSDK, but then of course there’s no way to actually test that your changes to the nxsdk_board do what you want if you can’t run it.

With respect to the above, the input to the single compartment neurons would be the output of the individual neurons of the previous Ensemble (i.e. the previous Conv layer). Now, those input compartments would be connected in pairs to another compartment which is supposed to be the immediate parent of them, whose (i.e. the parent’s) output would then be connected to the individual neurons of the next Ensemble.

So of course, unless I have understood things wrongly, there will be just one Connection from the previous Ensemble to my layer of compartmental neurons, and just one Connection again from the output (of my layer of compartmental neurons) to the next Ensemble, but the connections are supposed to be neuron-wise. Will this neuron-wise connections between Ensembles’ neurons and compartmental neurons be possible with NengoLoihi? I can see that we can obtain the input/output indices of the compartment neurons of the ens block (you created earlier) as below:

and I believe, the following creates the one to one connection from the inp Node and out Node to the compartmental neurons of the ens block:

but this is with Node to Ensemble (and Ensemble to Node).

Also, I guess… I will have to take help of NxSDK APIs (and not NengoLoihi) to enforce the child - parent relationship of the compartment neurons in my custom layer. Right?

And thanks for the confirmation of the requirement of a physical board to test out the custom layer of compartmental neurons, but I am slightly confused by the following:

Does it mean that the board and nxsdk_board objects can still be created in absence of a physical board? Here, I am learning something different, hence the confusion: " If the board isn’t available, the nxsdk_board object cannot be created (or used).".

It’s really a hack to be creating the nxsdk_board without hardware being available. All I’m saying is nothing in the code checks that hardware is available until with sim:, so theoretically anything before that can run as long as NxSDK is installed.

Yes.

Yes. In NengoLoihi, neurons are currently equivalent to compartments. i.e. each neuron is a single compartment neuron. So if the neurons you want to create are 3-compartment neurons, then you can create an Ensemble with 3 * n neurons. How you index things is up to you; e.g. you could have your inputs be the first 2 * n compartments, and use the last n compartments for outputs. Then you’d connect your inputs to ens.neurons[:2*n], and connect your outputs from ens.neurons[2*n:]. Then you just have to use NxSDK to wire things up appropriately within the Ensemble.

The above resolves my doubts for now @Eric. I will now wait to make some progress with the code before asking informed questions. This discussion has pretty much given me some ideas about heading into that direction.