The MLIR Transform Dialect. Your compiler is more powerful than you think

TL;DR

The paper tackles the challenge of insufficient, coarse-grained compiler control in heterogeneous hardware contexts by introducing the MLIR Transform dialect, an IR-based transformation language that lets performance engineers compose and steer existing compiler features without rebuilding the compiler. It delivers an interpretable Transform IR, handles, and parameters, plus static and dynamic verification through pre-/post-conditions and IRDL-based checks, enabling robust, reusable transformation pipelines. Five case studies demonstrate expressiveness, low overhead (up to $2.6\%$ compile-time), and practical benefits such as debugging optimization patterns, fine-grained performance tuning, and autotuning-driven exploration. The work significantly broadens access to domain knowledge in compilation, enables tight integration with optimization search methods, and supports scalable, portable, and high-performance code generation across dialects and targets.

Abstract

To take full advantage of a specific hardware target, performance engineers need to gain control on compilers in order to leverage their domain knowledge about the program and hardware. Yet, modern compilers are poorly controlled, usually by configuring a sequence of coarse-grained monolithic black-box passes, or by means of predefined compiler annotations/pragmas. These can be effective, but often do not let users precisely optimize their varying compute loads. As a consequence, performance engineers have to resort to implementing custom passes for a specific optimization heuristic, requiring compiler engineering expert knowledge. In this paper, we present a technique that provides fine-grained control of general-purpose compilers by introducing the Transform dialect, a controllable IR-based transformation system implemented in MLIR. The Transform dialect empowers performance engineers to optimize their various compute loads by composing and reusing existing - but currently hidden - compiler features without the need to implement new passes or even rebuilding the compiler. We demonstrate in five case studies that the Transform dialect enables precise, safe composition of compiler transformations and allows for straightforward integration with state-of-the-art search methods.

Figures (11)

• Figure 1: A Transform script performing code hoisting, loop splitting, tiling, and unrolling is defined (a) and used to transform the initial payload IR (b) to the transformed IR (c). The operations associated with specific handles are colored similarly.
• Figure 2: Pre-/post-conditions of convert_scf_to_cf declare which kinds of payload operations are consumed and removed (all operations from the scf dialect), and which new operations are introduced by this transform (the operations listed explicitly in lines 2 and 3).
• Figure 3: IRDL definition of the memref.subview operation. Highlighted parts are added in a copy of definition that expresses the post-condition of the transform.expand_strided_metadata transformation.
• Figure 4: constraints on input and output handles of transform.expand_strided_metadata
• Figure 5: Three alternative options for when to perform the automatic differentiation transformation. We use a transform to introspect the pipeline and automatically infer the appropriate transformation option.