Pre-compile Functions with `cond`.
Closed this issue · 3 comments
Hi all,
I think that we should pre-compile (using jit(inline=True)) functions containing cond. Once compiled this will make no difference to the performance (see #168), but it will improve the outside of compilation performance since cond recompiles the branches every time it is run. I believe that this is the source of the long runtimes that we have been experiencing using code.
Regards
Jordan.
A second motivation for this is that it will make the profiling more accurately reflect the performance that we are expecting.
So it turns out that this is not easy to do in practice so I am going to abandon this.
@benjaminpope will be familiar with the pitfalls of compiling class methods which include the requirement to mark self as static. I expected equinox to be a get out of jail free card with regards to this, but it was not. In fact because the module could be flattened it actually meant that any class with a device array leaf (all of them) could not be marked static at all 😬.