Reductions within neuron kernels

[neworderofjamie]

Being able to perform reductions across neuron kernels would be very useful e.g. for calculating softmax values across the output layer of a classifier. This is currently annoying to implement in GeNN except by adding an additional layers of neurons to calculate the sum / reducing on the host - both of which add one or more timesteps of latency.

Implementing this sort of reduction in CUDA would 'traditionally' require shared memory and liberal application of __syncthreads but, for populations of < warp size output neurons (which, in the context of output classes, is not unreasonable), the following algorithm can be used to provide each neuron with the sum of some value held in a register across a N neuron population:

for (int i=1; i<N; i*=2) {
    value += __shfl_xor_sync((1 << N) - 1, value, i);
}

This reduces following this pattern (adds value from thread i to the value from the thread with increasing bit swapped):

We could code generate unrolled versions of the above which could be inserted into code like this:

const scalar expY = exp($(Y));
scalar sumExpY = $(reduce_add, expY);
$(Pi) = expPi / sumExpPi;

Min and max could also be implemented in a similar way to allow more numerically stable softmaxes to also be implemented and this approach could also be extended populations of up to block size by performing additional reductions via shared memory. However, it would still be rather a leaky abstraction as, not only is it dependent on warp/block size, but only works outside of any conditional code. Furthermore, there would be no obvious implementation for the single-threaded CPU backend. Nonetheless, I think it would be useful. Thoughts?

[tnowotny]

Wouldn't it just have to be a simple loop on the single threaded CPU backend? That's simple, there is nothing more efficient as far as I can tell - so no harm done?

[tnowotny]

BTW, interesting pattern - additions are not done in the same order in each thread, so could one get in principle different final results in different threads due to rounding? (I admit, probably not of that much practical relevance)

[neworderofjamie]

But, in single threaded CPU mode, only one neuron is being updated at a time so while it's obviously trivial to calculate softmax I don't know how you would do it within the neuron code as that relies on a warps worth of neurons being updated simultaneously.

[neworderofjamie]

After writing it out, I think you're right that due to non-associativity different threads could have different values

[kernfel]

If different values are a problem, you could always add a single value = __shfl_sync(mask, value, 0). Speaking of mask: iirc, all threads in the warp have to be converged and participate in the call, even if N<32... something to be aware of.

[neworderofjamie]

Indeed although that form of reduction only ends with one thread having the result (which can admittedly be easily rectified using a final __shfl_sync). I think the newer 'sync' forms of these intrinsics wait for all threads to be converged but, issues like that are exactly my concern about this being a leaky abstraction (also realised that the xor-based reduction only works with power-of-two < warp size populations which is getting more and more niche...

[tnowotny]

But, in single threaded CPU mode, only one neuron is being updated at a time so while it's obviously trivial to calculate softmax I don't know how you would do it within the neuron code as that relies on a warps worth of neurons being updated simultaneously.

... good point ... it couldn't really be in the neuron code (and the generated loop over all neurons in the single threaded CPU code) in the traditional way. One would probably need to generate a separate code bit (loop) to do the reduction (a function with an "execute once" mechanism, so that the reduction is done once, all other neurons just receive the value?).

[neworderofjamie]

Well, this has spawned a much more interesting discussion than expected 😄 I think due to the caveats we've collectively come up with, the GeNN feature I original proposed would not be a good idea - for now one to save for CUDA-specific, case-specific neuron model hackery (such as where I originally applied it https://github.com/neworderofjamie/genn_examples/blob/batch_learning/eprop/evidence_accumulation/model.cc#L40-L43)

[tnowotny]

fair enough :-)