Presburger Arithmetic in MLIR

5 minute read

Published:

Presburger arithmetic provides the mathematical core for the polyhedral compilation techniques that drive analytical cache models, loop optimization for ML and HPC, formal verification, and hardware design. Arjun Pitchanathan, Kunwar Shaanjeet Singh Grover, Michel Weber and Tobias Grosser develop FPL, a fast new Presburger library, worked with the LLVM community to make Presburger arithmetic an integral part of the production-focused MLIR compiler framework.

With FPL now a native, performant, and approachable component of the LLVM ecosystem, compiler developers can co-develop the Presburger library and polyhedral transformations with ease. Expect that more scientists and compiler engineers will be able to jointly develop even more IMPACTful polyhedral optimizations.

Introduction

Polyhedral compilation based on Presburger arithmetic is used for performance optimization in high-performance computing and machine learning, formal verification, cache modeling, the derivation of data movement bounds, and configurable computing.

Over the last thirty years, several libraries provided mathematical foundations for polyhedral compilation. Omega and Polylib, developed many of the early ideas and inspired the design of polyhedral math libraries. The Parma Polyhedral Library (PPL) was the first polyhedral library to be used in GCC, as part of the Graphite; isl further improved on its robustness, provided the foundations for LLVM/Polly and Tensor Comprehensions, among others, and is considered state-of-the-art in polyhedral libraries. Finally, VPL demonstrated formally verified implementations for certain operations on sets of rationals.

As a Presburger library for LLVM/Polly, GCC, and many other tools, isl outgrew its mathematical roots quickly and increasingly incorporated broader compiler concepts. In particular, isl offered code generation facilities as well as higher-level code representations such as schedule trees. While schedule trees made polyhedral compilers more powerful, they were incompatible with the design and needs of LLVM. MLIR then generalized schedule trees with nested regions, and used these with the Affine dialect to create a simplified loop model that integrates well with the broader LLVM ecosystem. Until now it was necessary to choose between a well-integrated simple loop optimizer and full polyhedral transformations that translate MLIR dialects into isl-compatible formats. Unfortunately, the conversion costs and community disconnects hindered deeper integration of such approaches into MLIR.

With a generalization of schedule trees available in MLIR, the last missing piece is a fast and compute library for Presburger arithmetic. We fill in this gap by contributing a Presburger library, FPL, to LLVM. FPL was written from the ground up to meet the community’s standards, incentivizing further participation in the development of these mathematical foundations and allowing for the polyhedral compiler to be co-developed with the Presburger library.

Using FPL with MLIR

FPL provides four main data structures for Presburger arithmetic:

  • IntegerRelation
  • IntegerPolyhedron
  • PresburgerSet
  • PresburgerRelation

IntegerRelation and IntegerPolyhedron are relations and sets respectively defined by a set of affine constraints that must all be satisfied by every element in the set or relation. They support existentially quantified variables, so they can be thought of as projections of convex objects. Both can have symbols, i.e. parameters that have constant but unknown values, such as the size of an array that does not change during the course of a loop but whose is not known at compile-time.

The classes PresburgerSet and PresburgerRelations are unions of IntegerPolyhedrons and IntegerRelations respectively, they can support disjunctions of constraints.

We support integer-exact set operations like union, intersect, subtract, complement, equality checks, emptiness checks, and lexicographic optimization. More algorithmic details on these operation, see Pitchanathan et al. or the MLIR documentation.

MLIR API

MLIR provides support to define domain-specific IRs known as dialects. One example is the Affine dialect, which represents affine loop nests in a format suitable for polyhedral compilation. We present a example of the kind of analyses FPL enables over the Affine dialect. In this example, we wish to check whether the load in the second loop depends on the store of the first one.

In this case, we are checkint whether it is permissible to fuse the two loops on the left. This is allowed if there is no memory access dependency between an iteration of the i loop with an earlier iteration of the j loop. We can easily check this by converting the MemRefAccess objects created from the load/store IR operations to the Presburger library’s IntegerRelations and performing some Presburger operations.

In MLIR, a MemRefAccess represents a memory load or store access. We take the MemRefAccesses built from the store/load instructions and call getAccessRelation to obtain an IntegerRelation mapping the loop index to the memory location accessed. Inverting the second relation gives us a relation from memory locations accessed by the second loop to its induction variable. Composing these gives a relation representing the dependencies between the iterations i of the first loop to the iterations j of the second. By adding a constraint i >= j+1, we obtain a relation that has the dependencies from iterations of the first loop to earlier iterations of the second. These two loops can be fused only if there are no such dependencies, which we can check bu running the integer emptiness check on this final dependence relation. In this case, there exists a dependence between j = 1 and i = 2, so fusing the loops would be a miscompile. By providing support for Presburger operations, we lay the foundations for future work leveraging FPL for more complex polyhedral compilation techniques.

Use Example