Adding `compile_function` as execute option.

cubed-dev / cubed

Bounded-memory serverless distributed N-dimensional array processing

https://cubed-dev.github.io/cubed/

Apache License 2.0

122 stars 14 forks source link

Adding `compile_function` as execute option. #536

Closed alxmrs closed 3 months ago

alxmrs commented 3 months ago

This allows users to use jit or aot compilation during the dag finalization process within executors. It should work straightforwardly on jax/numba style jit compilation. It's possible, but maybe ugly, to perform jax-aot-style compilation.

This should help make progress towards #304.

(Note to reviewer: I realized this could be merged ahead of #508, so I pulled it out into a non-dependent PR).

alxmrs commented 3 months ago

I think this PR is ready to merge. The next things to explore in this area are:

[ ] Testing AOT compilation
[ ] Figuring out if this is the interface we need for exporting compiled functions (https://github.com/cubed-dev/cubed/issues/304#issuecomment-2258602859)

tomwhite commented 3 months ago

Thanks @alxmrs!