Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

slap jit from numba on it #2

Open
8bignic8 opened this issue Jul 11, 2023 · 1 comment
Open

slap jit from numba on it #2

8bignic8 opened this issue Jul 11, 2023 · 1 comment

Comments

@8bignic8
Copy link

resource: https://numba.pydata.org/numba-doc/latest/user/5minguide.html
from numba import njit

@georgedorn
Copy link

georgedorn commented Jul 11, 2023

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants