Lambda the Ultimate SSA: Optimizing Functional Programs in SSA

TL;DR

This paper introduces a novel SSA extension called regions to better optimize functional programs, enabling classical SSA reasoning for higher-order constructs and demonstrating its effectiveness in the LEAN4 theorem prover.

Contribution

We extend classical SSA to include regions, allowing systematic optimization of functional programs within an SSA framework, and implement a feature-complete backend for LEAN4.

Findings

Region optimizations achieve performance parity with existing LEAN4 backend.

The SSA+regions framework effectively models functional sub-expressions.

Our approach unifies optimization techniques for functional and imperative languages.

Abstract

Static Single Assignment (SSA) is the workhorse of modern optimizing compilers for imperative programming languages. However, functional languages have been slow to adopt SSA and prefer to use intermediate representations based on minimal lambda calculi due to SSA's inability to express higher order constructs. We exploit a new SSA construct -- regions -- in order to express functional optimizations via classical SSA based reasoning. Region optimization currently relies on ad-hoc analyses and transformations on imperative programs. These ad-hoc transformations are sufficient for imperative languages as regions are used in a limited fashion. In contrast, we use regions pervasively to model sub-expressions in our functional IR. This motivates us to systematize region optimizations. We extend classical SSA reasoning to regions for functional-style analyses and transformations. We implement a new SSA+regions based backend for LEAN4, a theorem prover that implements a purely functional, dependently typed programming language. Our backend is feature-complete and handles all constructs of LEAN4's functional intermediate representation {\lambda}rc within the SSA framework. We evaluate our proposed region optimizations by optimizing {\lambda}rc within an SSA+regions based framework implemented in MLIR and demonstrating performance parity with the current LEAN4 backend. We believe our work will pave the way for a unified optimization framework capable of representing, analyzing, and optimizing both functional and imperative languages.