JAX – NumPy on the CPU, GPU, and TPU, with great automatic differentiation

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by shihab on Friday September 29, 2023 @02:05AM (#37699836)

> 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.

An interesting article I read recently that has some relevance to this discussion. https://pharr.org/matt/blog/2018/04/18/ispc-origins

[ Parent | on HN ]

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by fulafel on Friday September 29, 2023 @02:27AM (#37699987)

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.

[ Parent | on HN ]

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by incrudible on Friday September 29, 2023 @06:19AM (#37701577)

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.

[ Parent | on HN ]