Re:Automated Code Optimization with E-Graphs
by jpsamaroo on Monday January 03, 2022 @09:36AM
(#29781148)
I think the OP's point is that running these optimizations in production as-is is dangerous because future code changes in the various places could accidentally impede the optimizer's ability to apply all the transformations that users of the codebase expect.
The obvious solution is to query the optimizer to get the final transformation as actual Julia code, and replace the pre-transform code with the post-transform optimized code, and disable any further optimization (aside from very trivial transforms that aren't worth directly including). This ensures that one doesn't accidentally lose the amazing benefits of this symbolic optimization approach on a given piece of code, and that production code always keeps its performance and correctness.
Another application we have in mind is to do similar linear algebra transformations to XLA. For example, sequential matrix-vector operations (A*v1 + A*v2) can be more optimally be applied as a single BLAS3 matrix-matrix call (A*[v1;v2]). Those kinds of rules can be very easily structured as an E-graph, and doing it this way would make it very inviting to mathematicians to extend and maintain the rulesets.
Both of those are cases where you may not want those passes to always be running, but they are super helpful in many numerical applications. And there's a lot more where that came from. Forward-mode AD can be implemented quite naturally this way. And there's a few more complex examples I'll hold in my back pocket until they are completed.