3 ms·

When NumPy is too slow

26 points by ingve 3y ago

credit_guy 3y ago
I don't get it. If code written in numpy is slow because the algorithm is bad, and you need to rewrite it, we can't say "numpy is too slow". A good chunk of the blogpost talks about this particular case, which is a non-case. The rest is some vague suggestions to use numba, or jax. But if your code is vectorizable in my experience you don't get any benefit from numba and jax. With the exception that jax is by default single-precision so you get some speed up from that. If you make it double precision, you get back to numpy speed. I suppose the author's conclusion is not all that different, otherwise they'd put a code snippet and show some numbers.
- tnecniv 3y ago
  I have definitely gotten significant speed ups using Jax. I try to vectorize my code as much as possible but that’s not always feasible for every aspect of an algorithm. Even for the vectorized operations, the JIT compilation gets me a speed up when used correctly (I.e. structuring code to minimize recompiling). Jax also has the added bonus of trivially running your algorithm on the GPU.
- akasakahakada 3y ago
  > default single-precision Did't know about that. So basically attach .astype(np.float16) to any native numpy array can have the same speed up without any existence of jax.
  - credit_guy 3y ago
    I think float16 is called half precision. Jax uses as default float32 and numpy uses float64. I didn't check how fast numpy is if you force the calculations to be a different precision. I'll give it a try and come back.
    - credit_guy 3y ago
      As promised, I did the experiment. Converting all floats in numpy to float32 achieves a speedup of about 20%. Writing the same code in jax (with the few small changes needed for example for random number generation), the speedup is of a factor of 2. But once we switch back to float64, jax becomes a tiny bit slower than numpy, maybe about 2% slower. Adding compilation to jax does not help if the code is already fully vectorized.
mathisfun123 3y ago
These articles are kind of funny - they cargo cult things that are generally sound so you can't really be mad but it still reads like a shibboleth. Like the claim that jax's jit speeds things up or that it looks at whole functions. Fact: jax doesn't have a jit on its own, at least not a machine code jit (maybe they cache ir/lowerings). Jax has an IR and a ~semi-novel perspective on tracing. That's it. Note that tracing means basically the exact opposite of "whole function" compilation. Jax also has two backends XLA and LLVM (through MLIR/stablehlo). Now I'll honestly plead ignorance of XLA - it probably does have some kind of jit but what it definitely has is hyperoptimized kernel implementations (thanks to hundreds of millions in dev time invested by Google). LLVM definitely has a jit but is it really that good that it beats numpy which itself is pretty well optimized (since it's so mature)? Then the suggestion to use numba: numba does do whole function compilation but numba just punts to the LLVM jit. So like what's significance of suggesting one vs the other? So this is my point: it's always a good thing when people search for optimization opportunities (more perf, less wall time, less power consumption, etc etc everyone wins) but it's not enough to just buy into the buzzwords. Or at least one shouldn't buy in and then present themselves as having actually done the due diligence. In serious compiler/perf engineering work everything is measured and remeasured and ablated and no one will take your seriously if you don't present your kernel counters (I'm kidding on the one but not really). Not everyone needs to be a compiler engineer but then you shouldn't pretend either.
akasakahakada 3y ago
numpy.sort is always the bottleneck of any numpy code. I won't suggest numba. That never worked on me, it never stop breaking and complaining with nonsensial error message. If I want faster I rather write fortran or nim then compile it to python, instead of broken numba.

26 points by ingve 3y ago

I don't get it. If code written in numpy is slow because the algorithm is bad, and you need to rewrite it, we can't say "numpy is too slow". A good chunk of the blogpost talks about this particular case, which is a non-case. The rest is some vague suggestions to use numba, or jax. But if your code is vectorizable in my experience you don't get any benefit from numba and jax. With the exception that jax is by default single-precision so you get some speed up from that. If you make it double precision, you get back to numpy speed. I suppose the author's conclusion is not all that different, otherwise they'd put a code snippet and show some numbers.

tnecniv 3y ago
I have definitely gotten significant speed ups using Jax. I try to vectorize my code as much as possible but that’s not always feasible for every aspect of an algorithm. Even for the vectorized operations, the JIT compilation gets me a speed up when used correctly (I.e. structuring code to minimize recompiling). Jax also has the added bonus of trivially running your algorithm on the GPU.
akasakahakada 3y ago
> default single-precision Did't know about that. So basically attach .astype(np.float16) to any native numpy array can have the same speed up without any existence of jax.
- credit_guy 3y ago
  I think float16 is called half precision. Jax uses as default float32 and numpy uses float64. I didn't check how fast numpy is if you force the calculations to be a different precision. I'll give it a try and come back.
  - credit_guy 3y ago
    As promised, I did the experiment. Converting all floats in numpy to float32 achieves a speedup of about 20%. Writing the same code in jax (with the few small changes needed for example for random number generation), the speedup is of a factor of 2. But once we switch back to float64, jax becomes a tiny bit slower than numpy, maybe about 2% slower. Adding compilation to jax does not help if the code is already fully vectorized.

mathisfun123 3y ago

These articles are kind of funny - they cargo cult things that are generally sound so you can't really be mad but it still reads like a shibboleth. Like the claim that jax's jit speeds things up or that it looks at whole functions. Fact: jax doesn't have a jit on its own, at least not a machine code jit (maybe they cache ir/lowerings). Jax has an IR and a ~semi-novel perspective on tracing. That's it. Note that tracing means basically the exact opposite of "whole function" compilation. Jax also has two backends XLA and LLVM (through MLIR/stablehlo). Now I'll honestly plead ignorance of XLA - it probably does have some kind of jit but what it definitely has is hyperoptimized kernel implementations (thanks to hundreds of millions in dev time invested by Google). LLVM definitely has a jit but is it really that good that it beats numpy which itself is pretty well optimized (since it's so mature)? Then the suggestion to use numba: numba does do whole function compilation but numba just punts to the LLVM jit. So like what's significance of suggesting one vs the other? So this is my point: it's always a good thing when people search for optimization opportunities (more perf, less wall time, less power consumption, etc etc everyone wins) but it's not enough to just buy into the buzzwords. Or at least one shouldn't buy in and then present themselves as having actually done the due diligence. In serious compiler/perf engineering work everything is measured and remeasured and ablated and no one will take your seriously if you don't present your kernel counters (I'm kidding on the one but not really). Not everyone needs to be a compiler engineer but then you shouldn't pretend either.

akasakahakada 3y ago

numpy.sort is always the bottleneck of any numpy code. I won't suggest numba. That never worked on me, it never stop breaking and complaining with nonsensial error message. If I want faster I rather write fortran or nim then compile it to python, instead of broken numba.