Equality Saturation for Optimizing High-Level Julia IR

2025-02-24

AI Summary

This work addresses the challenge that general-purpose compiler optimizations cannot effectively incorporate domain-specific expert knowledge. We propose a domain-customized optimization framework targeting Julia’s high-level intermediate representation (IR), built upon e-graphs and equality saturation. The framework enables domain experts to encode their knowledge via type-constrained rewrite rules. It introduces, for the first time in Julia’s compilation pipeline, control-flow graph (CFG) skeleton relaxation and integer linear programming (ILP)-driven, dominance-aware e-graph extraction—enabling unified rewriting of both pure and effectful functions. Our method fully models Julia’s multiple dispatch, control-flow, and side-effect semantics. Experimental evaluation demonstrates significant performance improvements on domain-specific kernels and low-level IR constructs such as broadcasting, while preserving semantic correctness and enabling quantitative assessment of compilation overhead.

Technology Category

• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Search and Optimization: Heuristic Search
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Vertical and domain-specific search

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 emph{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.

Problem

• Optimize Julia code using rewrite rules
• Apply optimizations with control flow
• Enable domain-specific code rewrites efficiently

Innovation

• Uses e-graphs for optimization
• Applies rewrite rules in Julia
• Incorporates ILP for e-graph extraction