mikepound/cubes

slap jit from numba on it

8bignic8 opened this issue · 1 comments

georgedorn commented

Some quick findings from some naive njit-slapping:

  • it may be possible, but it's nowhere near as simple as applying @njit (or @jit where that fails). the code will need some restructuring to isolate functionality into numpy-only sections vs python-containing sections.
  • in particular, quite a few numpy methods are not implemented in numba. array.pad() and np.packbits() (in a PR) are noteworthy examples.
  • numba doesn't support the yield keyword, so all_rotations() will need to be rewritten to return an array of copies of the cubes, rather than doing in-place modifications of the cubes. The speedups of numba will need to exceed the speedups of doing in-place modifications rather than copying the cube 24 times. ETA: Removing yield from this and applying @jit makes no speed difference. @njit may help, but it balks at np.rot90() so another implementation would need to be found.
  • applying @njit to crop_cube() and @jit to the packbits() version of rle(), generate_polycubes() and expand_cube() does help. At n=9, there's a time reduction from 70.83s to 48.748s. As n increases, the reduction should only improve, as the compilation time is a fixed amount of overhead for each run.

This may be worthwhile, but is easier said than done.

ETA: A minimal-effort implementation of this change is: https://github.com/georgedorn/cubes/tree/trial/numba This isn't particularly ready to go, but it does have the speedup from above and produces correct results. If somebody wants to take this an run with it, I suspect there are some deep optimizations to be done yet by doing the work to replace @jit with @njit.