the issue here is that if your ideal algorithm isn't simply expressible in numpy (which many aren't), you're pretty much out of luck. As a result, imo the better approach is to use a fast language that also compiles to GPU (e.g. Julia)
Having used JAX quite a bit for numerical computing (and having lectured on this use-case) I would say that a surprisingly large number of algorithms can be expressed as array[0] operations (even if it sometimes takes a bit of thinking).
And, more importantly, things that cannot be expressed that way tend to not be a good fit for GPU computing anyway (independently of the language / framework you are using).
[0]: `array` is a shortcut here, JAX is not limited to operations on arrays.
Agreed. I've done a fair amount of reworking signal processing algorithms to run on GPU/TPU, and it's a different beast. You often have to really rebuild the algorithm from the ground up to take advantage of parallelization. But often you /can/ rework the algorithm, and end up with much higher throughput than the crusty old serial algorithm: there's typically nothing fundamentally stopping you from finding a good implementation, just that the original devs were working in the 70s and hasn't thought that far ahead.
> things that cannot be expressed that way tend to not be a good fit for GPU computing anyway
I'll have to disagree with you a little bit here. SIMT model of GPUs are quiet a bit more expressive than the numpy's SIMD model. As an obvious example, you'll have to manually maintain a mask to implement if/else i.e. code path divergence in SIMD. GPUs automatically does this and many more to make your life easier. And frankly, I find it lot more easier to reason about what should happen to one data point than a bunch of them together.
Convenient authoring doesn't necessarily make it a good fit for the hardware. Add in enough divergence and your GPU code is going to be matched or outperformed by a competent CPU implementation (on a chip of comparable size). Branchless code can result in substantial speedups on either.
To be fair though, modern GPUs are pretty good at branching and latency hiding, while numpy-style code has poor data locality unless you have a magic compiler.
To spell out what the linked ISPC post implies, most of the difference, like ISPC shows, is differences in GPU languages and compilers vs CPU side equivalents.
Pretty sure I got multiple 1000x speed ups when I vectorized my my algo trader from a dumb python loop to a dumb numba compiled thing, and when I benchmarked Jax, the performance blew away the numba thing (which was already a million times faster than the naive version) because Jax performance stayed perfectly flat as the scale went up whereas numba slowed down. Might have been my approach for each, but it was enlightening and funny to watch.
I have at least one complaint with the numpy model:
When you chain a sequence of vectorized operations on arrays, loop fusion would save you from allocating memory for each intermediate variable, and the round trip time of moving it from RAM to CPU multiple times. I don’t know how good JAX’s JITted loop fusion is on CPU, but I’ve been very very impressed by Julia.
Eg: I had some Numpy code that took hours (and needed terabyte RAM) that was very straightforward to code in Julia, and needed only a few GB to finish in a few seconds — on my laptop.
I want to be able to think in arrays, but to also not have to materialize the arrays as much as possible.
I found this to be not true in practice when working with graphs.
Having access to high performance explicit loops and ifs/masks allows one to focus on the hard parts of the algorithms, rather than on the purely incidental puzzle how to best avoid spending time in the Python runtime.
An alternative is to write most of the program in Python + JAX + implement a few custom XLA ops in CUDA / Triton. That way, the program is very readable and can interoperate with the larger ecosystem, while still being fast to run.
A key difference is that each iteration of scan is called by the host. Put differently, JAX can't fuse scan into a single GPU kernel, but launches a kernel for each iteration.
Depening on the workload this is no problem. If you have many cheap iterations, you will notice the overhead.
I am not sure if they are working on fusing scan and what's the current status.
Yes, but is it really the same? Afaik the `unroll=n` parameter translates `n` iterations into a vanilla `for` loop which is then unrolled into sequential statements (in contrast to a JAX `fori` loop). There still is no loop on the accelerator, strictly speaking?
I think this is up to XLA to handle not Jax. The whole selling point in TF of the tf.function decorator (which uses XLA underneath as well) is that it fuses arithmetic to lower launch count.
