Equality Saturation for Optimizing High-Level Julia IR

TL;DR

This paper introduces an equality saturation-based system for optimizing high-level Julia code, enabling domain-specific rewrites with control flow, side effects, and type constraints, improving compiler effectiveness.

Contribution

It presents a novel system leveraging e-graphs and equality saturation for domain-specific Julia code optimization, including ILP-based e-graph extraction and CFG relaxation techniques.

Findings

Effective rewrites on high-level Julia code demonstrated

Supports control flow and side effects in optimizations

Analyzes compilation time for practical use

Abstract

Compilers are indispensable for transforming code written in high-level languages into performant machine code, but their general-purpose optimizations sometimes fall short. Domain experts might be aware of certain optimizations that the compiler is unable to apply or that are only valid in a particular domain. We have developed a system that allows domain experts to express rewrite rules to optimize code in the Julia programming language. Our system builds on e-graphs and equality saturation. It can apply optimizations in the presence of control flow and side effects. As Julia uses multiple dispatch, we allow users to constrain rewrite rules by argument types, and propagate type information through the e-graph representation. We propose an ILP formulation for optimal e-graph extraction taking into account dominance properties for code reuse and introduce CFG skeleton relaxation to rewrite calls to pure functions as well as those with side effects. Use cases demonstrate that our system can perform rewrites on high-level, domain-specific code, as well as on lower-level code such as Julia's broadcasting mechanism. Finally, we analyze the required compilation time.