These commits wrap up the previous Softmax refactor. All the important changes are on the CUDA side. Once I unified the code I also added a special instantiation that made the kernels faster in certain cases (small inner dim, large softmax dim - might be useful in NLP for short sequences?)

tl;dr CUDA Softmax now supports a dim argument, and is usually 4x-256x faster than the previous implementation (it didn't have a spatial implementation before). Now, it also shares kernels with LogSoftmax, and certain optimizations benefited the log case giving up to 64x speedup in certain cases as well.

Here are the plots that show old / new timing ratios for different sizes of dim and size of the innermost dimensions (on the left of dim). Paralellizing over batch is easy, so it is fixed as 64 in all plots. Red dots are better, blue are regressions. Note that the plot is log in all axis (so z of 8 means 2^8x faster)

Softmax

Benefits from this diff all over the place. The old kernel was written in a quite archaic way.

LogSoftmax

Benefits from adding a custom kernel for the cases when inner_size is no longer 1, so we can't use the super fast kernel, but the dim_size is large, so using a single thread to reduce values is slow. It is only enabled for a subset of the space where it provided speedups.

Overall times

These are the log plots (in all axes) of the time (no more ratios) for the new algorithm. Softmax on the left, LogSoftmax on the right: