Precise control of compilers: a practical approach to principled optimization

Abstract

Production compilers remain constrained by rigid designs and inflexible optimization strategies that limit their adaptation to evolving hardware designs and application requirements. While academic research continues to produce innovative compilation techniques and intermediate representations that enable more flexibility and control over the compilation process, bridging the gap between these theoretical advances and practical implementation remains a significant challenge. The emergence of the MLIR compiler infrastructure, with its extensible dialect-based architecture, presents a unique opportunity to address two fundamental challenges in modern compiler design: how to integrate principled representations of programs and transformations into production compilers and how to provide users with meaningful control over optimization decisions. For the representation challenge, we first integrate Rise, a functional pattern-based intermediate representation, as an MLIR dialect to establish the foundation for systematic rewrite-based optimizations in a production-ready compiler infrastructure. We demonstrate the practical utility of this integration by leveraging Rise's pattern-based abstractions to compile a machine learning model. While Rise's functional foundation provides natural guarantees about program properties during transformations, these are lost when integrating with MLIR's largely unconstrained rewriting system. Transformations do not guarantee the preservation of dialect-specific invariants or even the Static Single Assignment (SSA) form beyond dynamic checks. We present an approach for statically validating that a set of MLIR rewrites maintains these critical properties.

For the control challenge, we introduce two distinct approaches to transformation control. The Transform dialect, now part of upstream MLIR, encodes optimization sequences as an MLIR dialect. It exposes previously hidden compiler optimizations, enabling developers to express complex transformation strategies directly in MLIR's infrastructure and moving beyond rigid pass pipelines to provide precise control over which optimizations are applied to specific parts of the program.

To provide a principled foundation for composing complex optimization strategies, we adapt Elevate, a functional language for composing transformations, to the MLIR ecosystem. We extend Elevate with MLIR-specific constructs while making MLIR's IR immutable to enable both approaches to work in unison. This adaptation enables compiler developers to express sophisticated optimization strategies simply by composing rewrite rules. We show how this approach scales to flexibly match computational structures, such as attention, in machine learning models.

By systematically addressing both representation and control challenges, we establish principled foundations for optimization while maintaining precise control, demonstrating how modern compiler infrastructures can evolve beyond rigid designs without sacrificing practical applicability.